remove sharding that has been upstreamed into mlx-lm

LLama, Ministral, GPT-OSS, MiniMax and DeepSeek all have their tensor
sharding upstream. We can remove the sharding code from exo.

.github/workflows/build-app.yml

@@ -1,5 +1,16 @@
 name: Build EXO macOS DMG
 
+# Release workflow:
+# 1. Create a draft GitHub Release with the tag name (e.g. v1.0.0) and write release notes in markdown
+# 2. Push the tag: git tag v1.0.0 && git push origin v1.0.0
+# 3. This workflow builds, signs, and notarizes the DMG
+# 4. Release notes are embedded in appcast.xml for Sparkle (rendered as markdown)
+# 5. DMG and appcast.xml are uploaded to S3
+# 6. The draft GitHub Release is published with the DMG attached
+#
+# For alpha releases (e.g. v1.0.0-alpha.1): draft release and notes are optional.
+# If no draft exists, a release is auto-created with generated notes.
+
 on:
   workflow_dispatch:
   push:
@@ -11,8 +22,10 @@ on:
 jobs:
   build-macos-app:
     runs-on: "macos-26"
+    permissions:
+      contents: write
     env:
-      SPARKLE_VERSION: 2.8.1
+      SPARKLE_VERSION: 2.9.0-beta.1
       SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }}
       SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }}
       SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }}
@@ -87,6 +100,52 @@ jobs:
             exit 1
           fi
 
+      - name: Fetch and validate release notes
+        if: github.ref_type == 'tag'
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          # Find draft release by name using gh release list (more reliable with default token)
+          echo "Looking for draft release named '$GITHUB_REF_NAME'..."
+          DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "")
+
+          if [[ -z "$DRAFT_EXISTS" ]]; then
+            if [[ "$IS_ALPHA" == "true" ]]; then
+              echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)"
+              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
+              exit 0
+            fi
+            echo "ERROR: No draft release found for tag $GITHUB_REF_NAME"
+            echo "Please create a draft release with release notes before pushing the tag."
+            exit 1
+          fi
+
+          # Fetch full release details via API to get body and ID
+          echo "Found draft release, fetching details..."
+          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "")
+
+          # Extract release notes
+          NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""')
+          if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then
+            if [[ "$IS_ALPHA" == "true" ]]; then
+              echo "Draft release has no notes (optional for alphas)"
+              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
+              exit 0
+            fi
+            echo "ERROR: Draft release exists but has no release notes"
+            echo "Please add release notes to the draft release before pushing the tag."
+            exit 1
+          fi
+
+          # Save release ID for later publishing
+          RELEASE_ID=$(echo "$RELEASE_JSON" | jq -r '.id')
+          echo "DRAFT_RELEASE_ID=$RELEASE_ID" >> $GITHUB_ENV
+          echo "HAS_RELEASE_NOTES=true" >> $GITHUB_ENV
+
+          echo "Found draft release (ID: $RELEASE_ID), saving release notes..."
+          echo "$NOTES" > /tmp/release_notes.md
+          echo "RELEASE_NOTES_FILE=/tmp/release_notes.md" >> $GITHUB_ENV
+
       # ============================================================
       # Install dependencies
       # ============================================================
@@ -304,6 +363,28 @@ jobs:
             $CHANNEL_FLAG \
             .
 
+      - name: Inject release notes into appcast
+        if: github.ref_type == 'tag' && env.HAS_RELEASE_NOTES == 'true'
+        env:
+          RELEASE_VERSION: ${{ env.RELEASE_VERSION }}
+        run: |
+          # Inject markdown release notes with sparkle:format="markdown" (Sparkle 2.9+)
+          export NOTES=$(cat "$RELEASE_NOTES_FILE")
+
+          # Insert description after the enclosure tag for this version
+          awk '
+            /<enclosure[^>]*>/ && index($0, ENVIRON["RELEASE_VERSION"]) {
+              print
+              print "            <description sparkle:format=\"markdown\"><![CDATA["
+              print ENVIRON["NOTES"]
+              print "            ]]></description>"
+              next
+            }
+            { print }
+          ' output/appcast.xml > output/appcast.xml.tmp && mv output/appcast.xml.tmp output/appcast.xml
+
+          echo "Injected markdown release notes for version $RELEASE_VERSION"
+
       # ============================================================
       # Upload artifacts
       # ============================================================
@@ -336,3 +417,26 @@ jobs:
             aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"
             aws s3 cp appcast.xml "s3://${SPARKLE_S3_BUCKET}/${PREFIX}appcast.xml" --content-type application/xml --cache-control no-cache
           fi
+
+      - name: Publish GitHub Release
+        if: github.ref_type == 'tag'
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          DMG_PATH="output/EXO-${RELEASE_VERSION}.dmg"
+
+          if [[ "$HAS_RELEASE_NOTES" == "true" ]]; then
+            # Update the draft release with the tag and upload DMG
+            gh api --method PATCH "repos/${{ github.repository }}/releases/$DRAFT_RELEASE_ID" \
+              -f tag_name="$GITHUB_REF_NAME" \
+              -F draft=false
+            gh release upload "$GITHUB_REF_NAME" "$DMG_PATH" --clobber
+            echo "Published release $GITHUB_REF_NAME with DMG attached"
+          else
+            # Alpha without draft release - create one with auto-generated notes
+            gh release create "$GITHUB_REF_NAME" "$DMG_PATH" \
+              --title "$GITHUB_REF_NAME" \
+              --generate-notes \
+              --prerelease
+            echo "Created alpha release $GITHUB_REF_NAME with auto-generated notes"
+          fi

.mlx_typings/mlx_lm/generate.pyi

@@ -276,23 +276,24 @@ class BatchGenerator:
         logprobs: mx.array
         finish_reason: Optional[str]
 
-    unprocessed_prompts: List[Any]
-
     def __init__(
         self,
-        model: nn.Module,
+        model,
         max_tokens: int = ...,
-        stop_tokens: Optional[set[int]] = ...,
+        stop_tokens: Optional[set] = ...,
         sampler: Optional[Callable[[mx.array], mx.array]] = ...,
         completion_batch_size: int = ...,
         prefill_batch_size: int = ...,
         prefill_step_size: int = ...,
     ) -> None: ...
     def insert(
-        self, prompts: List[List[int]], max_tokens: Union[List[int], int, None] = ...
-    ) -> List[int]: ...
-    def stats(self) -> BatchStats: ...
-    def next(self) -> List[Response]: ...
+        self, prompts, max_tokens: Union[List[int], int, None] = ...
+    ):  # -> list[Any]:
+        ...
+    def stats(self):  # -> BatchStats:
+        ...
+    def next(self):  # -> list[Any]:
+        ...
 
 def batch_generate(
     model,

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -39,18 +39,12 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    tokens: list[int]
-    def reset(self) -> None: ...
-    def add_token(self, token: int) -> None: ...
-    def finalize(self) -> None: ...
+    def reset(self): ...
+    def add_token(self, token): ...
+    def finalize(self): ...
     @property
-    def text(self) -> str:
-        """The full text decoded so far."""
-        ...
-    @property
-    def last_segment(self) -> str:
+    def last_segment(self):
         """Return the last segment of readable text since last time this property was accessed."""
-        ...
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
     """NaiveStreamingDetokenizer relies on the underlying tokenizer
@@ -114,7 +108,6 @@ class TokenizerWrapper:
     _tokenizer: PreTrainedTokenizerFast
     eos_token_id: int | None
     eos_token: str | None
-    eos_token_ids: list[int] | None
     bos_token_id: int | None
     bos_token: str | None
     vocab_size: int

AGENTS.md

@@ -40,6 +40,31 @@ uv run ruff check
 nix fmt
 ```
 
+## Pre-Commit Checks (REQUIRED)
+
+**IMPORTANT: Always run these checks before committing code. CI will fail if these don't pass.**
+
+```bash
+# 1. Type checking - MUST pass with 0 errors
+uv run basedpyright
+
+# 2. Linting - MUST pass
+uv run ruff check
+
+# 3. Formatting - MUST be applied
+nix fmt
+
+# 4. Tests - MUST pass
+uv run pytest
+```
+
+Run all checks in sequence:
+```bash
+uv run basedpyright && uv run ruff check && nix fmt && uv run pytest
+```
+
+If `nix fmt` changes any files, stage them before committing. The CI runs `nix flake check` which verifies formatting, linting, and runs Rust tests.
+
 ## Architecture
 
 ### Node Composition
@@ -91,45 +116,6 @@ From .cursorrules:
 - Catch exceptions only where you can handle them meaningfully
 - Use `@final` and immutability wherever applicable
 
-## Model Storage
-
-Downloaded models are stored in `~/.exo/models/` (not the standard HuggingFace cache location).
-
-## Creating Model Instances via API
-
-When testing with the API, you must first create a model instance before sending chat completions:
-
-```bash
-# 1. Get instance previews for a model
-curl "http://localhost:52415/instance/previews?model_id=llama-3.2-1b"
-
-# 2. Create an instance from the first valid preview
-INSTANCE=$(curl -s "http://localhost:52415/instance/previews?model_id=llama-3.2-1b" | jq -c '.previews[] | select(.error == null) | .instance' | head -n1)
-curl -X POST http://localhost:52415/instance -H 'Content-Type: application/json' -d "{\"instance\": $INSTANCE}"
-
-# 3. Wait for the runner to become ready (check logs for "runner ready")
-
-# 4. Send chat completions using the full model ID
-curl -X POST http://localhost:52415/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{"model": "mlx-community/Llama-3.2-1B-Instruct-4bit", "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 50}'
-```
-
-## Logs
-
-Exo logs are stored in `~/.exo/exo.log`. This is useful for debugging runner crashes and distributed issues.
-
 ## Testing
 
 Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests.
-
-### Distributed Testing
-
-When running distributed tests across multiple machines, use `EXO_LIBP2P_NAMESPACE` to isolate your test cluster from other exo instances on the same network:
-
-```bash
-# On each machine in the test cluster, use the same unique namespace
-EXO_LIBP2P_NAMESPACE=my-test-cluster uv run exo
-```
-
-This prevents your test cluster from discovering and interfering with production or other developers' exo clusters.

README.md

@@ -27,13 +27,22 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
 - **Tensor Parallelism**: exo supports sharding models, for up to 1.8x speedup on 2 devices and 3.2x speedup on 4 devices.
 - **MLX Support**: exo uses [MLX](https://github.com/ml-explore/mlx) as an inference backend and [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) for distributed communication.
 
+## Dashboard
+
+exo includes a built-in dashboard for managing your cluster and chatting with models.
+
+<p align="center">
+  <img src="docs/imgs/dashboard-cluster-view.png" alt="exo dashboard - cluster view showing 4 x M3 Ultra Mac Studio with DeepSeek v3.1 and Kimi-K2-Thinking loaded" width="80%" />
+</p>
+<p align="center"><em>4 × 512GB M3 Ultra Mac Studio running DeepSeek v3.1 (8-bit) and Kimi-K2-Thinking (4-bit)</em></p>
+
 ## Benchmarks
 
 <details>
   <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>
 
@@ -41,7 +50,7 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
   <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>
 
@@ -49,7 +58,7 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
   <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>
 
@@ -154,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.
@@ -166,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.
@@ -312,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.
@@ -320,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.

TODO.md

@@ -19,6 +19,7 @@
 25. Rethink retry logic
 26. Task cancellation. When API http request gets cancelled, it should cancel corresponding task.
 27. Log cleanup - per-module log filters and default to DEBUG log levels
+28. Validate RDMA connections with ibv_devinfo in the info gatherer
 
 Potential refactors:
 

app/EXO/EXO.xcodeproj/project.pbxproj

@@ -585,7 +585,7 @@
 			repositoryURL = "https://github.com/sparkle-project/Sparkle.git";
 			requirement = {
 				kind = upToNextMajorVersion;
-				minimumVersion = 2.8.1;
+				minimumVersion = 2.9.0-beta.1;
 			};
 		};
 /* End XCRemoteSwiftPackageReference section */

app/EXO/EXO.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -6,8 +6,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/sparkle-project/Sparkle.git",
       "state" : {
-        "revision" : "5581748cef2bae787496fe6d61139aebe0a451f6",
-        "version" : "2.8.1"
+        "revision" : "e641adb41915a8409895e2e30666aa64e487b637",
+        "version" : "2.9.0-beta.1"
       }
     }
   ],

app/EXO/EXO/ContentView.swift

@@ -56,6 +56,11 @@ struct ContentView: View {
     }
 
     private var shouldShowLocalNetworkWarning: Bool {
+        // Show warning if local network is not working and EXO is running.
+        // The checker uses a longer timeout on first launch to allow time for
+        // the permission prompt, so this correctly handles both:
+        // 1. User denied permission on first launch
+        // 2. Permission broke after restart (macOS TCC bug)
         if case .notWorking = localNetworkChecker.status {
             return controller.status != .stopped
         }

app/EXO/EXO/Services/LocalNetworkChecker.swift

@@ -5,8 +5,8 @@ import os.log
 /// Checks if the app's local network permission is actually functional.
 ///
 /// macOS local network permission can appear enabled in System Preferences but not
-/// actually work after a restart. This service detects this by creating a UDP
-/// connection to the mDNS multicast address (224.0.0.251:5353).
+/// actually work after a restart. This service uses NWConnection to mDNS multicast
+/// to verify actual connectivity.
 @MainActor
 final class LocalNetworkChecker: ObservableObject {
     enum Status: Equatable {
@@ -35,30 +35,43 @@ final class LocalNetworkChecker: ObservableObject {
     }
 
     private static let logger = Logger(subsystem: "io.exo.EXO", category: "LocalNetworkChecker")
+    private static let hasCompletedInitialCheckKey = "LocalNetworkChecker.hasCompletedInitialCheck"
 
     @Published private(set) var status: Status = .unknown
-    @Published private(set) var lastConnectionState: String = "none"
 
     private var connection: NWConnection?
     private var checkTask: Task<Void, Never>?
 
+    /// Whether we've completed at least one check (stored in UserDefaults)
+    private var hasCompletedInitialCheck: Bool {
+        get { UserDefaults.standard.bool(forKey: Self.hasCompletedInitialCheckKey) }
+        set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) }
+    }
+
     /// Checks if local network access is working.
     func check() {
         checkTask?.cancel()
         status = .checking
-        lastConnectionState = "connecting"
+
+        // Use longer timeout on first launch to allow time for permission prompt
+        let isFirstCheck = !hasCompletedInitialCheck
+        let timeout: UInt64 = isFirstCheck ? 30_000_000_000 : 3_000_000_000
 
         checkTask = Task { [weak self] in
             guard let self else { return }
-            let result = await self.performCheck()
+
+            Self.logger.info("Checking local network connectivity (first check: \(isFirstCheck))")
+            let result = await self.checkConnectivity(timeout: timeout)
             self.status = result
+            self.hasCompletedInitialCheck = true
+
             Self.logger.info("Local network check complete: \(result.displayText)")
         }
     }
 
-    private func performCheck() async -> Status {
-        Self.logger.info("Checking local network access via UDP multicast")
-
+    /// Checks connectivity using NWConnection to mDNS multicast.
+    /// The connection attempt triggers the permission prompt if not yet shown.
+    private func checkConnectivity(timeout: UInt64) async -> Status {
         connection?.cancel()
         connection = nil
 
@@ -84,22 +97,7 @@ final class LocalNetworkChecker: ObservableObject {
                 continuation.resume(returning: status)
             }
 
-            conn.stateUpdateHandler = { [weak self] state in
-                let stateStr: String
-                switch state {
-                case .setup: stateStr = "setup"
-                case .preparing: stateStr = "preparing"
-                case .ready: stateStr = "ready"
-                case .waiting(let e): stateStr = "waiting(\(e))"
-                case .failed(let e): stateStr = "failed(\(e))"
-                case .cancelled: stateStr = "cancelled"
-                @unknown default: stateStr = "unknown"
-                }
-
-                Task { @MainActor in
-                    self?.lastConnectionState = stateStr
-                }
-
+            conn.stateUpdateHandler = { state in
                 switch state {
                 case .ready:
                     resumeOnce(.working)
@@ -108,6 +106,7 @@ final class LocalNetworkChecker: ObservableObject {
                     if errorStr.contains("54") || errorStr.contains("ECONNRESET") {
                         resumeOnce(.notWorking(reason: "Connection blocked"))
                     }
+                // Otherwise keep waiting - might be showing permission prompt
                 case .failed(let error):
                     let errorStr = "\(error)"
                     if errorStr.contains("65") || errorStr.contains("EHOSTUNREACH")
@@ -127,7 +126,7 @@ final class LocalNetworkChecker: ObservableObject {
             conn.start(queue: .main)
 
             Task {
-                try? await Task.sleep(nanoseconds: 3_000_000_000)
+                try? await Task.sleep(nanoseconds: timeout)
                 let state = conn.state
                 switch state {
                 case .ready:

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -6,7 +6,7 @@ enum NetworkSetupHelper {
     private static let logger = Logger(subsystem: "io.exo.EXO", category: "NetworkSetup")
     private static let daemonLabel = "io.exo.networksetup"
     private static let scriptDestination =
-        "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
+        "/Library/Application Support/EXO/disable_bridge.sh"
     private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist"
     private static let requiredStartInterval: Int = 1791
 
@@ -28,35 +28,6 @@ enum NetworkSetupHelper {
         # 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
@@ -141,6 +112,13 @@ enum NetworkSetupHelper {
         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(

bench/exo_bench.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import argparse
+import contextlib
 import http.client
 import json
 import os
@@ -15,9 +16,6 @@ from urllib.parse import urlencode
 from loguru import logger
 from transformers import AutoTokenizer
 
-from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.types.memory import Memory
-
 
 class ExoHttpError(RuntimeError):
     def __init__(self, status: int, reason: str, body_preview: str):
@@ -26,7 +24,7 @@ class ExoHttpError(RuntimeError):
 
 
 class ExoClient:
-    def __init__(self, host: str, port: int, timeout_s: float = 2400.0):
+    def __init__(self, host: str, port: int, timeout_s: float = 600.0):
         self.host = host
         self.port = port
         self.timeout_s = timeout_s
@@ -104,22 +102,46 @@ def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
 
+def runner_failed(runner: dict[str, Any]) -> bool:
+    return "RunnerFailed" in runner
+
+
+def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
+    if "RunnerFailed" in runner:
+        return runner["RunnerFailed"].get("errorMessage")
+    return None
+
+
 def wait_for_instance_ready(
     client: ExoClient, instance_id: str, timeout: float = 24000.0
 ) -> None:
     start_time = time.time()
+    instance_existed = False
     while time.time() - start_time < timeout:
         state = client.request_json("GET", "/state")
         instances = state.get("instances", {})
 
         if instance_id not in instances:
+            if instance_existed:
+                # Instance was deleted after being created - likely due to runner failure
+                raise RuntimeError(
+                    f"Instance {instance_id} was deleted (runner may have failed)"
+                )
             time.sleep(0.1)
             continue
 
+        instance_existed = True
         instance = instances[instance_id]
         runner_ids = runner_ids_from_instance(instance)
         runners = state.get("runners", {})
 
+        # Check for failed runners first
+        for rid in runner_ids:
+            runner = runners.get(rid, {})
+            if runner_failed(runner):
+                error_msg = get_runner_failed_message(runner) or "Unknown error"
+                raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+
         if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
             return
 
@@ -241,6 +263,9 @@ class PromptSizer:
             ids = tokenizer.apply_chat_template(
                 messages, tokenize=True, add_generation_prompt=True
             )
+            # Fix for transformers 5.x
+            if hasattr(ids, "input_ids"):
+                ids = ids.input_ids
             return int(len(ids))
 
         return count_fn
@@ -296,6 +321,12 @@ def main() -> int:
         default=4,
         help="Only consider placements using <= this many nodes.",
     )
+    ap.add_argument(
+        "--min-nodes",
+        type=int,
+        default=1,
+        help="Only consider placements using >= this many nodes.",
+    )
     ap.add_argument(
         "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
     )
@@ -317,7 +348,7 @@ def main() -> int:
         help="Warmup runs per placement (uses first pp/tg).",
     )
     ap.add_argument(
-        "--timeout", type=float, default=2400.0, help="HTTP timeout (seconds)."
+        "--timeout", type=float, default=600.0, help="HTTP timeout (seconds)."
     )
     ap.add_argument(
         "--json-out",
@@ -396,7 +427,7 @@ def main() -> int:
         ):
             continue
 
-        if 0 < n <= args.max_nodes:
+        if args.min_nodes <= n <= args.max_nodes:
             selected.append(p)
 
     if not selected:
@@ -438,7 +469,13 @@ def main() -> int:
         )
 
         client.request_json("POST", "/instance", body={"instance": instance})
-        wait_for_instance_ready(client, instance_id)
+        try:
+            wait_for_instance_ready(client, instance_id)
+        except (RuntimeError, TimeoutError) as e:
+            logger.error(f"Failed to initialize placement: {e}")
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            continue
 
         time.sleep(1)
 
@@ -450,17 +487,17 @@ def main() -> int:
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
             for pp in pp_list:
-                if (
-                    pp * n_nodes > 2048
-                    and "ring" in instance_meta.lower()
-                    and "tensor" in sharding.lower()
-                ):
-                    model_card = MODEL_CARDS[short_id]
-                    if model_card.metadata.storage_size > Memory.from_gb(10):
-                        logger.info(
-                            f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
-                        )
-                        continue
+                # if (
+                #     pp * n_nodes > 2048
+                #     and "ring" in instance_meta.lower()
+                #     and "tensor" in sharding.lower()
+                # ):
+                #     model_card = MODEL_CARDS[short_id]
+                #     if model_card.metadata.storage_size > Memory.from_gb(10):
+                #         logger.info(
+                #             f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
+                #         )
+                #         continue
                 for tg in tg_list:
                     runs: list[dict[str, Any]] = []
                     for r in range(args.repeat):

dashboard/package-lock.json

@@ -863,7 +863,6 @@
 			"integrity": "sha512-oH8tXw7EZnie8FdOWYrF7Yn4IKrqTFHhXvl8YxXxbKwTMcD/5NNCryUSEXRk2ZR4ojnub0P8rNrsVGHXWqIDtA==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@standard-schema/spec": "^1.0.0",
 				"@sveltejs/acorn-typescript": "^1.0.5",
@@ -903,7 +902,6 @@
 			"integrity": "sha512-Y1Cs7hhTc+a5E9Va/xwKlAJoariQyHY+5zBgCZg4PFWNYQ1nMN9sjK1zhw1gK69DuqVP++sht/1GZg1aRwmAXQ==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@sveltejs/vite-plugin-svelte-inspector": "^4.0.1",
 				"debug": "^4.4.1",
@@ -1520,7 +1518,6 @@
 			"integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"undici-types": "~6.21.0"
 			}
@@ -1530,7 +1527,6 @@
 			"resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz",
 			"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
 			"license": "MIT",
-			"peer": true,
 			"bin": {
 				"acorn": "bin/acorn"
 			},
@@ -1943,7 +1939,6 @@
 			"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
 			"dev": true,
 			"license": "ISC",
-			"peer": true,
 			"engines": {
 				"node": ">=12"
 			}
@@ -2651,7 +2646,6 @@
 			"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"engines": {
 				"node": ">=12"
 			},
@@ -2839,7 +2833,6 @@
 			"resolved": "https://registry.npmjs.org/svelte/-/svelte-5.45.3.tgz",
 			"integrity": "sha512-ngKXNhNvwPzF43QqEhDOue7TQTrG09em1sd4HBxVF0Wr2gopAmdEWan+rgbdgK4fhBtSOTJO8bYU4chUG7VXZQ==",
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@jridgewell/remapping": "^2.3.4",
 				"@jridgewell/sourcemap-codec": "^1.5.0",
@@ -2984,7 +2977,6 @@
 			"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
 			"dev": true,
 			"license": "Apache-2.0",
-			"peer": true,
 			"bin": {
 				"tsc": "bin/tsc",
 				"tsserver": "bin/tsserver"
@@ -3006,7 +2998,6 @@
 			"integrity": "sha512-+Oxm7q9hDoLMyJOYfUYBuHQo+dkAloi33apOPP56pzj+vsdJDzr+j1NISE5pyaAuKL4A3UD34qd0lx5+kfKp2g==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"esbuild": "^0.25.0",
 				"fdir": "^6.4.4",

dashboard/src/lib/components/MarkdownContent.svelte

@@ -53,62 +53,285 @@
 	marked.use({ renderer });
 
 	/**
-	 * Preprocess LaTeX: convert \(...\) to $...$ and \[...\] to $$...$$
-	 * Also protect code blocks from LaTeX processing
+	 * Unescape HTML entities that marked may have escaped
+	 */
+	function unescapeHtmlEntities(text: string): string {
+		return text
+			.replace(/&lt;/g, '<')
+			.replace(/&gt;/g, '>')
+			.replace(/&amp;/g, '&')
+			.replace(/&quot;/g, '"')
+			.replace(/&#39;/g, "'");
+	}
+
+	// Storage for math expressions extracted before markdown processing
+	const mathExpressions: Map<string, { content: string; displayMode: boolean }> = new Map();
+	let mathCounter = 0;
+
+	// Storage for HTML snippets that need protection from markdown
+	const htmlSnippets: Map<string, string> = new Map();
+	let htmlCounter = 0;
+
+	// Use alphanumeric placeholders that won't be interpreted as HTML tags
+	const MATH_PLACEHOLDER_PREFIX = 'MATHPLACEHOLDER';
+	const CODE_PLACEHOLDER_PREFIX = 'CODEPLACEHOLDER';
+	const HTML_PLACEHOLDER_PREFIX = 'HTMLPLACEHOLDER';
+
+	/**
+	 * Preprocess LaTeX: extract math, handle LaTeX document commands, and protect content
 	 */
 	function preprocessLaTeX(text: string): string {
-		// Protect code blocks
+		// Reset storage
+		mathExpressions.clear();
+		mathCounter = 0;
+		htmlSnippets.clear();
+		htmlCounter = 0;
+
+		// Protect code blocks first
 		const codeBlocks: string[] = [];
 		let processed = text.replace(/```[\s\S]*?```|`[^`]+`/g, (match) => {
 			codeBlocks.push(match);
-			return `<<CODE_${codeBlocks.length - 1}>>`;
+			return `${CODE_PLACEHOLDER_PREFIX}${codeBlocks.length - 1}END`;
 		});
 
-		// Convert \(...\) to $...$
-		processed = processed.replace(/\\\((.+?)\\\)/g, '$$$1$');
-		
-		// Convert \[...\] to $$...$$
-		processed = processed.replace(/\\\[([\s\S]*?)\\\]/g, '$$$$$1$$$$');
+		// Remove LaTeX document commands
+		processed = processed.replace(/\\documentclass(\[[^\]]*\])?\{[^}]*\}/g, '');
+		processed = processed.replace(/\\usepackage(\[[^\]]*\])?\{[^}]*\}/g, '');
+		processed = processed.replace(/\\begin\{document\}/g, '');
+		processed = processed.replace(/\\end\{document\}/g, '');
+		processed = processed.replace(/\\maketitle/g, '');
+		processed = processed.replace(/\\title\{[^}]*\}/g, '');
+		processed = processed.replace(/\\author\{[^}]*\}/g, '');
+		processed = processed.replace(/\\date\{[^}]*\}/g, '');
+
+		// Remove \require{...} commands (MathJax-specific, not supported by KaTeX)
+		processed = processed.replace(/\$\\require\{[^}]*\}\$/g, '');
+		processed = processed.replace(/\\require\{[^}]*\}/g, '');
+
+		// Remove unsupported LaTeX commands/environments (tikzpicture, figure, center, etc.)
+		processed = processed.replace(/\\begin\{tikzpicture\}[\s\S]*?\\end\{tikzpicture\}/g, () => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, '<div class="latex-diagram-placeholder"><span class="latex-diagram-icon">📐</span><span class="latex-diagram-text">Diagram</span></div>');
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\begin\{figure\}[\s\S]*?\\end\{figure\}/g, () => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, '<div class="latex-diagram-placeholder"><span class="latex-diagram-icon">🖼️</span><span class="latex-diagram-text">Figure</span></div>');
+			htmlCounter++;
+			return placeholder;
+		});
+		// Strip center environment (layout only, no content change)
+		processed = processed.replace(/\\begin\{center\}/g, '');
+		processed = processed.replace(/\\end\{center\}/g, '');
+		// Strip other layout environments
+		processed = processed.replace(/\\begin\{flushleft\}/g, '');
+		processed = processed.replace(/\\end\{flushleft\}/g, '');
+		processed = processed.replace(/\\begin\{flushright\}/g, '');
+		processed = processed.replace(/\\end\{flushright\}/g, '');
+		processed = processed.replace(/\\label\{[^}]*\}/g, '');
+		processed = processed.replace(/\\caption\{[^}]*\}/g, '');
+
+		// Protect escaped dollar signs (e.g., \$50 should become $50, not LaTeX)
+		processed = processed.replace(/\\\$/g, 'ESCAPEDDOLLARPLACEHOLDER');
+
+		// Convert LaTeX math environments to display math (both bare and wrapped in $...$)
+		const mathEnvs = ['align', 'align\\*', 'equation', 'equation\\*', 'gather', 'gather\\*', 'multline', 'multline\\*', 'eqnarray', 'eqnarray\\*', 'array', 'matrix', 'pmatrix', 'bmatrix', 'vmatrix', 'cases'];
+		for (const env of mathEnvs) {
+			// Handle $\begin{env}...\end{env}$ (with dollar signs, possibly multiline)
+			const wrappedRegex = new RegExp(`\\$\\\\begin\\{${env}\\}(\\{[^}]*\\})?([\\s\\S]*?)\\\\end\\{${env}\\}\\$`, 'g');
+			processed = processed.replace(wrappedRegex, (_, args, content) => {
+				const cleanEnv = env.replace('\\*', '*');
+				const mathContent = `\\begin{${cleanEnv}}${args || ''}${content}\\end{${cleanEnv}}`;
+				const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+				mathExpressions.set(placeholder, { content: mathContent, displayMode: true });
+				mathCounter++;
+				return placeholder;
+			});
+
+			// Handle bare \begin{env}...\end{env} (without dollar signs)
+			const bareRegex = new RegExp(`\\\\begin\\{${env}\\}(\\{[^}]*\\})?([\\s\\S]*?)\\\\end\\{${env}\\}`, 'g');
+			processed = processed.replace(bareRegex, (_, args, content) => {
+				const cleanEnv = env.replace('\\*', '*');
+				const mathContent = `\\begin{${cleanEnv}}${args || ''}${content}\\end{${cleanEnv}}`;
+				const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+				mathExpressions.set(placeholder, { content: mathContent, displayMode: true });
+				mathCounter++;
+				return placeholder;
+			});
+		}
+
+		// Convert LaTeX proof environments to styled blocks (use placeholders for HTML)
+		processed = processed.replace(
+			/\\begin\{proof\}([\s\S]*?)\\end\{proof\}/g,
+			(_, content) => {
+				const html = `<div class="latex-proof"><div class="latex-proof-header">Proof</div><div class="latex-proof-content">${content}</div></div>`;
+				const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+				htmlSnippets.set(placeholder, html);
+				htmlCounter++;
+				return placeholder;
+			}
+		);
+
+		// Convert LaTeX theorem-like environments
+		const theoremEnvs = ['theorem', 'lemma', 'corollary', 'proposition', 'definition', 'remark', 'example'];
+		for (const env of theoremEnvs) {
+			const envRegex = new RegExp(`\\\\begin\\{${env}\\}([\\s\\S]*?)\\\\end\\{${env}\\}`, 'gi');
+			const envName = env.charAt(0).toUpperCase() + env.slice(1);
+			processed = processed.replace(envRegex, (_, content) => {
+				const html = `<div class="latex-theorem"><div class="latex-theorem-header">${envName}</div><div class="latex-theorem-content">${content}</div></div>`;
+				const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+				htmlSnippets.set(placeholder, html);
+				htmlCounter++;
+				return placeholder;
+			});
+		}
+
+		// Convert LaTeX text formatting commands (use placeholders to protect from markdown)
+		processed = processed.replace(/\\emph\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<em>${content}</em>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\textit\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<em>${content}</em>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\textbf\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<strong>${content}</strong>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\texttt\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<code class="inline-code">${content}</code>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\underline\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<u>${content}</u>`);
+			htmlCounter++;
+			return placeholder;
+		});
+
+		// Handle LaTeX line breaks and spacing
+		processed = processed.replace(/\\\\(?:\s*\n)?/g, '\n'); // \\ -> newline
+		processed = processed.replace(/\\newline/g, '\n');
+		processed = processed.replace(/\\par\b/g, '\n\n');
+		processed = processed.replace(/\\quad/g, ' ');
+		processed = processed.replace(/\\qquad/g, '  ');
+		processed = processed.replace(/~~/g, ' '); // non-breaking space
+
+		// Remove other common LaTeX commands that don't render
+		processed = processed.replace(/\\centering/g, '');
+		processed = processed.replace(/\\noindent/g, '');
+		processed = processed.replace(/\\hfill/g, '');
+		processed = processed.replace(/\\vspace\{[^}]*\}/g, '');
+		processed = processed.replace(/\\hspace\{[^}]*\}/g, ' ');
+
+		// Convert \(...\) to placeholder (display: false)
+		processed = processed.replace(/\\\(([\s\S]+?)\\\)/g, (_, content) => {
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}INLINE${mathCounter}END`;
+			mathExpressions.set(placeholder, { content, displayMode: false });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Convert \[...\] to placeholder (display: true)
+		processed = processed.replace(/\\\[([\s\S]*?)\\\]/g, (_, content) => {
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+			mathExpressions.set(placeholder, { content, displayMode: true });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Extract display math ($$...$$) BEFORE markdown processing
+		processed = processed.replace(/\$\$([\s\S]*?)\$\$/g, (_, content) => {
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+			mathExpressions.set(placeholder, { content: content.trim(), displayMode: true });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Extract inline math ($...$) BEFORE markdown processing
+		// Allow single-line only, skip currency patterns like $5 or $50
+		processed = processed.replace(/\$([^\$\n]+?)\$/g, (match, content) => {
+			if (/^\d/.test(content.trim())) {
+				return match; // Keep as-is for currency
+			}
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}INLINE${mathCounter}END`;
+			mathExpressions.set(placeholder, { content: content.trim(), displayMode: false });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Restore escaped dollar signs
+		processed = processed.replace(/ESCAPEDDOLLARPLACEHOLDER/g, '$');
 
 		// Restore code blocks
-		processed = processed.replace(/<<CODE_(\d+)>>/g, (_, index) => codeBlocks[parseInt(index)]);
+		processed = processed.replace(new RegExp(`${CODE_PLACEHOLDER_PREFIX}(\\d+)END`, 'g'), (_, index) => codeBlocks[parseInt(index)]);
+
+		// Clean up any remaining stray backslashes from unrecognized commands
+		processed = processed.replace(/\\(?=[a-zA-Z])/g, ''); // Remove \ before letters (unrecognized commands)
 
 		return processed;
 	}
 
 	/**
-	 * Render math expressions with KaTeX after HTML is generated
+	 * Render math expressions with KaTeX and restore HTML placeholders
 	 */
 	function renderMath(html: string): string {
-		// Render display math ($$...$$)
-		html = html.replace(/\$\$([\s\S]*?)\$\$/g, (_, math) => {
-			try {
-				return katex.renderToString(math.trim(), {
-					displayMode: true,
-					throwOnError: false,
-					output: 'html'
-				});
-			} catch {
-				return `<span class="math-error">$$${math}$$</span>`;
-			}
-		});
+		// Replace all math placeholders with rendered KaTeX
+		for (const [placeholder, { content, displayMode }] of mathExpressions) {
+			const escapedPlaceholder = placeholder.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+			const regex = new RegExp(escapedPlaceholder, 'g');
 
-		// Render inline math ($...$) but avoid matching currency like $5
-		html = html.replace(/\$([^\$\n]+?)\$/g, (match, math) => {
-			// Skip if it looks like currency ($ followed by number)
-			if (/^\d/.test(math.trim())) {
-				return match;
-			}
-			try {
-				return katex.renderToString(math.trim(), {
-					displayMode: false,
-					throwOnError: false,
-					output: 'html'
-				});
-			} catch {
-				return `<span class="math-error">$${math}$</span>`;
-			}
-		});
+			html = html.replace(regex, () => {
+				try {
+					const rendered = katex.renderToString(content, {
+						displayMode,
+						throwOnError: false,
+						output: 'html'
+					});
+
+					if (displayMode) {
+						return `
+							<div class="math-display-wrapper">
+								<div class="math-display-header">
+									<span class="math-label">LaTeX</span>
+									<button type="button" class="copy-math-btn" data-math-source="${encodeURIComponent(content)}" title="Copy LaTeX source">
+										<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+											<rect width="14" height="14" x="8" y="8" rx="2" ry="2"/>
+											<path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/>
+										</svg>
+									</button>
+								</div>
+								<div class="math-display-content">
+									${rendered}
+								</div>
+							</div>
+						`;
+					} else {
+						return `<span class="math-inline">${rendered}</span>`;
+					}
+				} catch {
+					const display = displayMode ? `$$${content}$$` : `$${content}$`;
+					return `<span class="math-error"><span class="math-error-icon">⚠</span> ${display}</span>`;
+				}
+			});
+		}
+
+		// Restore HTML placeholders (for \textbf, \emph, etc.)
+		for (const [placeholder, htmlContent] of htmlSnippets) {
+			const escapedPlaceholder = placeholder.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+			const regex = new RegExp(escapedPlaceholder, 'g');
+			html = html.replace(regex, htmlContent);
+		}
 
 		return html;
 	}
@@ -154,16 +377,50 @@
 		}
 	}
 
+	async function handleMathCopyClick(event: Event) {
+		const target = event.currentTarget as HTMLButtonElement;
+		const encodedSource = target.getAttribute('data-math-source');
+		if (!encodedSource) return;
+
+		const source = decodeURIComponent(encodedSource);
+
+		try {
+			await navigator.clipboard.writeText(source);
+			// Show copied feedback
+			const originalHtml = target.innerHTML;
+			target.innerHTML = `
+				<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+					<path d="M20 6L9 17l-5-5"/>
+				</svg>
+			`;
+			target.classList.add('copied');
+			setTimeout(() => {
+				target.innerHTML = originalHtml;
+				target.classList.remove('copied');
+			}, 2000);
+		} catch (error) {
+			console.error('Failed to copy math:', error);
+		}
+	}
+
 	function setupCopyButtons() {
 		if (!containerRef || !browser) return;
 
-		const buttons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
-		for (const button of buttons) {
+		const codeButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
+		for (const button of codeButtons) {
 			if (button.dataset.listenerBound !== 'true') {
 				button.dataset.listenerBound = 'true';
 				button.addEventListener('click', handleCopyClick);
 			}
 		}
+
+		const mathButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-math-btn');
+		for (const button of mathButtons) {
+			if (button.dataset.listenerBound !== 'true') {
+				button.dataset.listenerBound = 'true';
+				button.addEventListener('click', handleMathCopyClick);
+			}
+		}
 	}
 
 	$effect(() => {
@@ -424,28 +681,290 @@
 		color: #60a5fa;
 	}
 
-	/* KaTeX math styling */
+	/* KaTeX math styling - Base */
 	.markdown-content :global(.katex) {
 		font-size: 1.1em;
+		color: oklch(0.9 0 0);
 	}
 
-	.markdown-content :global(.katex-display) {
+	/* Display math container wrapper */
+	.markdown-content :global(.math-display-wrapper) {
 		margin: 1rem 0;
+		border-radius: 0.5rem;
+		overflow: hidden;
+		border: 1px solid rgba(255, 215, 0, 0.15);
+		background: rgba(0, 0, 0, 0.3);
+		transition: border-color 0.2s ease, box-shadow 0.2s ease;
+	}
+
+	.markdown-content :global(.math-display-wrapper:hover) {
+		border-color: rgba(255, 215, 0, 0.25);
+		box-shadow: 0 0 12px rgba(255, 215, 0, 0.08);
+	}
+
+	/* Display math header - hidden by default, slides in on hover */
+	.markdown-content :global(.math-display-header) {
+		display: flex;
+		justify-content: space-between;
+		align-items: center;
+		padding: 0.375rem 0.75rem;
+		background: rgba(255, 215, 0, 0.03);
+		border-bottom: 1px solid rgba(255, 215, 0, 0.08);
+		opacity: 0;
+		max-height: 0;
+		padding-top: 0;
+		padding-bottom: 0;
+		overflow: hidden;
+		transition:
+			opacity 0.2s ease,
+			max-height 0.2s ease,
+			padding 0.2s ease;
+	}
+
+	.markdown-content :global(.math-display-wrapper:hover .math-display-header) {
+		opacity: 1;
+		max-height: 2.5rem;
+		padding: 0.375rem 0.75rem;
+	}
+
+	.markdown-content :global(.math-label) {
+		color: rgba(255, 215, 0, 0.7);
+		font-size: 0.65rem;
+		font-weight: 500;
+		text-transform: uppercase;
+		letter-spacing: 0.1em;
+		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
+	}
+
+	.markdown-content :global(.copy-math-btn) {
+		display: flex;
+		align-items: center;
+		justify-content: center;
+		padding: 0.25rem;
+		background: transparent;
+		border: none;
+		color: var(--exo-light-gray, #9ca3af);
+		cursor: pointer;
+		transition: color 0.2s;
+		border-radius: 0.25rem;
+		opacity: 0;
+		transition:
+			color 0.2s,
+			opacity 0.15s ease;
+	}
+
+	.markdown-content :global(.math-display-wrapper:hover .copy-math-btn) {
+		opacity: 1;
+	}
+
+	.markdown-content :global(.copy-math-btn:hover) {
+		color: var(--exo-yellow, #ffd700);
+	}
+
+	.markdown-content :global(.copy-math-btn.copied) {
+		color: #22c55e;
+	}
+
+	/* Display math content area */
+	.markdown-content :global(.math-display-content) {
+		padding: 1rem 1.25rem;
 		overflow-x: auto;
 		overflow-y: hidden;
-		padding: 0.5rem 0;
 	}
 
-	.markdown-content :global(.katex-display > .katex) {
+	/* Custom scrollbar for math overflow */
+	.markdown-content :global(.math-display-content::-webkit-scrollbar) {
+		height: 6px;
+	}
+
+	.markdown-content :global(.math-display-content::-webkit-scrollbar-track) {
+		background: rgba(255, 255, 255, 0.05);
+		border-radius: 3px;
+	}
+
+	.markdown-content :global(.math-display-content::-webkit-scrollbar-thumb) {
+		background: rgba(255, 215, 0, 0.2);
+		border-radius: 3px;
+	}
+
+	.markdown-content :global(.math-display-content::-webkit-scrollbar-thumb:hover) {
+		background: rgba(255, 215, 0, 0.35);
+	}
+
+	.markdown-content :global(.math-display-content .katex-display) {
+		margin: 0;
+		padding: 0;
+	}
+
+	.markdown-content :global(.math-display-content .katex-display > .katex) {
 		text-align: center;
 	}
 
+	/* Inline math wrapper */
+	.markdown-content :global(.math-inline) {
+		display: inline;
+		padding: 0 0.125rem;
+		border-radius: 0.25rem;
+		transition: background-color 0.15s ease;
+	}
+
+	.markdown-content :global(.math-inline:hover) {
+		background: rgba(255, 215, 0, 0.05);
+	}
+
+	/* Dark theme KaTeX overrides */
+	.markdown-content :global(.katex .mord),
+	.markdown-content :global(.katex .minner),
+	.markdown-content :global(.katex .mop),
+	.markdown-content :global(.katex .mbin),
+	.markdown-content :global(.katex .mrel),
+	.markdown-content :global(.katex .mpunct) {
+		color: oklch(0.9 0 0);
+	}
+
+	/* Fraction lines and rules */
+	.markdown-content :global(.katex .frac-line),
+	.markdown-content :global(.katex .overline-line),
+	.markdown-content :global(.katex .underline-line),
+	.markdown-content :global(.katex .hline),
+	.markdown-content :global(.katex .rule) {
+		border-color: oklch(0.85 0 0) !important;
+		background: oklch(0.85 0 0);
+	}
+
+	/* Square roots and SVG elements */
+	.markdown-content :global(.katex .sqrt-line) {
+		border-color: oklch(0.85 0 0) !important;
+	}
+
+	.markdown-content :global(.katex svg) {
+		fill: oklch(0.85 0 0);
+		stroke: oklch(0.85 0 0);
+	}
+
+	.markdown-content :global(.katex svg path) {
+		stroke: oklch(0.85 0 0);
+	}
+
+	/* Delimiters (parentheses, brackets, braces) */
+	.markdown-content :global(.katex .delimsizing),
+	.markdown-content :global(.katex .delim-size1),
+	.markdown-content :global(.katex .delim-size2),
+	.markdown-content :global(.katex .delim-size3),
+	.markdown-content :global(.katex .delim-size4),
+	.markdown-content :global(.katex .mopen),
+	.markdown-content :global(.katex .mclose) {
+		color: oklch(0.75 0 0);
+	}
+
+	/* Math error styling */
 	.markdown-content :global(.math-error) {
+		display: inline-flex;
+		align-items: center;
+		gap: 0.375rem;
 		color: #f87171;
 		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
 		font-size: 0.875em;
 		background: rgba(248, 113, 113, 0.1);
-		padding: 0.125rem 0.25rem;
+		padding: 0.25rem 0.5rem;
 		border-radius: 0.25rem;
+		border: 1px solid rgba(248, 113, 113, 0.2);
+	}
+
+	.markdown-content :global(.math-error-icon) {
+		font-size: 0.875em;
+		opacity: 0.9;
+	}
+
+	/* LaTeX proof environment */
+	.markdown-content :global(.latex-proof) {
+		margin: 1rem 0;
+		padding: 1rem 1.25rem;
+		background: rgba(255, 255, 255, 0.02);
+		border-left: 3px solid rgba(255, 215, 0, 0.4);
+		border-radius: 0 0.375rem 0.375rem 0;
+	}
+
+	.markdown-content :global(.latex-proof-header) {
+		font-weight: 600;
+		font-style: italic;
+		color: oklch(0.85 0 0);
+		margin-bottom: 0.5rem;
+	}
+
+	.markdown-content :global(.latex-proof-header::after) {
+		content: '.';
+	}
+
+	.markdown-content :global(.latex-proof-content) {
+		color: oklch(0.9 0 0);
+	}
+
+	.markdown-content :global(.latex-proof-content p:last-child) {
+		margin-bottom: 0;
+	}
+
+	/* QED symbol at end of proof */
+	.markdown-content :global(.latex-proof-content::after) {
+		content: '∎';
+		display: block;
+		text-align: right;
+		color: oklch(0.7 0 0);
+		margin-top: 0.5rem;
+	}
+
+	/* LaTeX theorem-like environments */
+	.markdown-content :global(.latex-theorem) {
+		margin: 1rem 0;
+		padding: 1rem 1.25rem;
+		background: rgba(255, 215, 0, 0.03);
+		border: 1px solid rgba(255, 215, 0, 0.15);
+		border-radius: 0.375rem;
+	}
+
+	.markdown-content :global(.latex-theorem-header) {
+		font-weight: 700;
+		color: var(--exo-yellow, #ffd700);
+		margin-bottom: 0.5rem;
+	}
+
+	.markdown-content :global(.latex-theorem-header::after) {
+		content: '.';
+	}
+
+	.markdown-content :global(.latex-theorem-content) {
+		color: oklch(0.9 0 0);
+		font-style: italic;
+	}
+
+	.markdown-content :global(.latex-theorem-content p:last-child) {
+		margin-bottom: 0;
+	}
+
+	/* LaTeX diagram/figure placeholder */
+	.markdown-content :global(.latex-diagram-placeholder) {
+		display: flex;
+		align-items: center;
+		justify-content: center;
+		gap: 0.5rem;
+		margin: 1rem 0;
+		padding: 1.5rem 2rem;
+		background: rgba(255, 255, 255, 0.02);
+		border: 1px dashed rgba(255, 215, 0, 0.25);
+		border-radius: 0.5rem;
+		color: rgba(255, 215, 0, 0.6);
+		font-size: 0.875rem;
+	}
+
+	.markdown-content :global(.latex-diagram-icon) {
+		font-size: 1.25rem;
+		opacity: 0.8;
+	}
+
+	.markdown-content :global(.latex-diagram-text) {
+		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
+		font-size: 0.75rem;
+		text-transform: uppercase;
+		letter-spacing: 0.05em;
 	}
 </style>

dashboard/src/lib/components/ModelCard.svelte

@@ -197,7 +197,7 @@ function toggleNodeDetails(nodeId: string): void {
 	// Uses API preview data when available, falls back to local estimation
 	const placementPreview = $derived(() => {
 		const nodeArray = nodeList();
-		if (nodeArray.length === 0) return { nodes: [], canFit: false, totalAvailable: 0, error: null };
+		if (nodeArray.length === 0) return { nodes: [], canFit: false, totalAvailable: 0, topoWidth: 260, topoHeight: 90, error: null };
 		
 		const numNodes = nodeArray.length;
 		const iconSize = numNodes === 1 ? 50 : 36;

dashboard/src/lib/components/TopologyGraph.svelte

@@ -1,7 +1,7 @@
 <script lang="ts">
 	import { onMount, onDestroy } from 'svelte';
 	import * as d3 from 'd3';
-import { topologyData, isTopologyMinimized, debugMode } from '$lib/stores/app.svelte';
+import { topologyData, isTopologyMinimized, debugMode, type NodeInfo } from '$lib/stores/app.svelte';
 
 	interface Props {
 		class?: string;
@@ -24,14 +24,14 @@ function getNodeLabel(nodeId: string): string {
 
 function getInterfaceLabel(nodeId: string, ip?: string): { label: string; missing: boolean } {
 	if (!ip) return { label: '?', missing: true };
-	
+
 	// Strip port if present (e.g., "192.168.1.1:8080" -> "192.168.1.1")
 	const cleanIp = ip.includes(':') && !ip.includes('[') ? ip.split(':')[0] : ip;
-	
+
 	// Helper to check a node's interfaces
-	function checkNode(node: typeof data.nodes[string]): string | null {
+	function checkNode(node: NodeInfo | undefined): string | null {
 		if (!node) return null;
-		
+
 		const matchFromInterfaces = node.network_interfaces?.find((iface) =>
 			(iface.addresses || []).some((addr) => addr === cleanIp || addr === ip)
 		);
@@ -39,17 +39,19 @@ function getInterfaceLabel(nodeId: string, ip?: string): { label: string; missin
 			return matchFromInterfaces.name;
 		}
 
-		const mapped = node.ip_to_interface?.[cleanIp] || node.ip_to_interface?.[ip];
-		if (mapped && mapped.trim().length > 0) {
-			return mapped;
+		if (node.ip_to_interface) {
+			const mapped = node.ip_to_interface[cleanIp] || (ip ? node.ip_to_interface[ip] : undefined);
+			if (mapped && mapped.trim().length > 0) {
+				return mapped;
+			}
 		}
 		return null;
 	}
-	
+
 	// Try specified node first
 	const result = checkNode(data?.nodes?.[nodeId]);
 	if (result) return { label: result, missing: false };
-	
+
 	// Fallback: search all nodes for this IP
 	for (const [, otherNode] of Object.entries(data?.nodes || {})) {
 		const otherResult = checkNode(otherNode);
@@ -255,21 +257,24 @@ function wrapLine(text: string, maxLen: number): string[] {
 		const arrowsGroup = svg.append('g').attr('class', 'arrows-group');
 		const debugLabelsGroup = svg.append('g').attr('class', 'debug-edge-labels');
 
-		const pairMap = new Map<string, { a: string; b: string; aToB: boolean; bToA: boolean; connections: Array<{ from: string; to: string; ip: string; ifaceLabel: string; missingIface: boolean }> }>();
-		let debugEdgeLabels: Array<{ connections: typeof pairMap extends Map<string, infer V> ? V['connections'] : never; isLeft: boolean; isTop: boolean; mx: number; my: number }> | null = null;
+		type ConnectionInfo = { from: string; to: string; ip: string; ifaceLabel: string; missingIface: boolean };
+		type PairEntry = { a: string; b: string; aToB: boolean; bToA: boolean; connections: ConnectionInfo[] };
+		type DebugEdgeLabelEntry = { connections: ConnectionInfo[]; isLeft: boolean; isTop: boolean; mx: number; my: number };
+		const pairMap = new Map<string, PairEntry>();
+		const debugEdgeLabels: DebugEdgeLabelEntry[] = [];
 		edges.forEach(edge => {
 			if (!edge.source || !edge.target || edge.source === edge.target) return;
 			if (!positionById[edge.source] || !positionById[edge.target]) return;
-			
+
 			const a = edge.source < edge.target ? edge.source : edge.target;
 			const b = edge.source < edge.target ? edge.target : edge.source;
 			const key = `${a}|${b}`;
 			const entry = pairMap.get(key) || { a, b, aToB: false, bToA: false, connections: [] };
-			
+
 			if (edge.source === a) entry.aToB = true;
 			else entry.bToA = true;
 
-			const ip = edge.sendBackIp || edge.sendBackMultiaddr?.ip_address || '?';
+			const ip = edge.sendBackIp || '?';
 			const ifaceInfo = getInterfaceLabel(edge.source, ip);
 			entry.connections.push({
 				from: edge.source,
@@ -338,9 +343,8 @@ function wrapLine(text: string, maxLen: number): string[] {
 				// Determine which side of viewport based on edge midpoint
 				const isLeft = mx < centerX;
 				const isTop = my < safeCenterY;
-				
+
 				// Store for batch rendering after all edges processed
-				if (!debugEdgeLabels) debugEdgeLabels = [];
 				debugEdgeLabels.push({
 					connections: entry.connections,
 					isLeft,
@@ -381,32 +385,32 @@ function wrapLine(text: string, maxLen: number): string[] {
 			}
 			
 			// Group by quadrant: topLeft, topRight, bottomLeft, bottomRight
-			const quadrants: Record<string, typeof debugEdgeLabels> = {
+			const quadrants: Record<string, DebugEdgeLabelEntry[]> = {
 				topLeft: [],
 				topRight: [],
 				bottomLeft: [],
 				bottomRight: []
 			};
-			
+
 			debugEdgeLabels.forEach(edge => {
 				const key = (edge.isTop ? 'top' : 'bottom') + (edge.isLeft ? 'Left' : 'Right');
 				quadrants[key].push(edge);
 			});
-			
+
 			// Render each quadrant
-			Object.entries(quadrants).forEach(([quadrant, edges]) => {
-				if (edges.length === 0) return;
-				
+			Object.entries(quadrants).forEach(([quadrant, quadrantEdges]) => {
+				if (quadrantEdges.length === 0) return;
+
 				const isLeft = quadrant.includes('Left');
 				const isTop = quadrant.includes('top');
-				
+
 				let baseX = isLeft ? padding : width - padding;
 				let baseY = isTop ? padding : height - padding;
 				const textAnchor = isLeft ? 'start' : 'end';
-				
+
 				let currentY = baseY;
-				
-				edges.forEach(edge => {
+
+				quadrantEdges.forEach(edge => {
 					edge.connections.forEach(conn => {
 						const arrow = getArrow(conn.from, conn.to);
 						const label = `${arrow} ${conn.ip} ${conn.ifaceLabel}`;

dashboard/src/lib/stores/app.svelte.ts

@@ -71,52 +71,66 @@ export interface Instance {
 	};
 }
 
-interface RawNodeProfile {
+// Granular node state types from the new state structure
+interface RawNodeIdentity {
 	modelId?: string;
 	chipId?: string;
 	friendlyName?: string;
-	networkInterfaces?: Array<{
-		name?: string;
-		ipAddress?: string;
-		addresses?: Array<{ address?: string } | string>;
-		ipv4?: string;
-		ipv6?: string;
-		ipAddresses?: string[];
-		ips?: string[];
-	}>;
-	memory?: {
-		ramTotal?: { inBytes: number };
-		ramAvailable?: { inBytes: number };
-		swapTotal?: { inBytes: number };
-		swapAvailable?: { inBytes: number };
-	};
-	system?: {
-		gpuUsage?: number;
-		temp?: number;
-		sysPower?: number;
+}
+
+interface RawMemoryUsage {
+	ramTotal?: { inBytes: number };
+	ramAvailable?: { inBytes: number };
+	swapTotal?: { inBytes: number };
+	swapAvailable?: { inBytes: number };
+}
+
+interface RawSystemPerformanceProfile {
+	gpuUsage?: number;
+	temp?: number;
+	sysPower?: number;
+	pcpuUsage?: number;
+	ecpuUsage?: number;
+}
+
+interface RawNetworkInterfaceInfo {
+	name?: string;
+	ipAddress?: string;
+	addresses?: Array<{ address?: string } | string>;
+	ipv4?: string;
+	ipv6?: string;
+	ipAddresses?: string[];
+	ips?: string[];
+}
+
+interface RawNodeNetworkInfo {
+	interfaces?: RawNetworkInterfaceInfo[];
+}
+
+interface RawSocketConnection {
+	sinkMultiaddr?: {
+		address?: string;
+		ip_address?: string;
+		address_type?: string;
+		port?: number;
 	};
 }
 
-interface RawTopologyNode {
-	nodeId: string;
-	nodeProfile: RawNodeProfile;
+interface RawRDMAConnection {
+	sourceRdmaIface?: string;
+	sinkRdmaIface?: string;
 }
 
-interface RawTopologyConnection {
-	localNodeId: string;
-	sendBackNodeId: string;
-	sendBackMultiaddr?:
-		| { multiaddr?: string; address?: string; ip_address?: string }
-		| string;
-}
+type RawConnectionEdge = RawSocketConnection | RawRDMAConnection;
+
+// New nested mapping format: { source: { sink: [edge1, edge2, ...] } }
+type RawConnectionsMap = Record<string, Record<string, RawConnectionEdge[]>>;
 
 interface RawTopology {
-	nodes: RawTopologyNode[];
-	connections?: RawTopologyConnection[];
+	nodes: string[];
+	connections?: RawConnectionsMap;
 }
 
-type RawNodeProfiles = Record<string, RawNodeProfile>;
-
 export interface DownloadProgress {
 	totalBytes: number;
 	downloadedBytes: number;
@@ -171,7 +185,11 @@ interface RawStateResponse {
 	>;
 	runners?: Record<string, unknown>;
 	downloads?: Record<string, unknown[]>;
-	nodeProfiles?: RawNodeProfiles;
+	// New granular node state fields
+	nodeIdentities?: Record<string, RawNodeIdentity>;
+	nodeMemory?: Record<string, RawMemoryUsage>;
+	nodeSystem?: Record<string, RawSystemPerformanceProfile>;
+	nodeNetwork?: Record<string, RawNodeNetworkInfo>;
 }
 
 export interface MessageAttachment {
@@ -206,56 +224,69 @@ export interface Conversation {
 
 const STORAGE_KEY = "exo-conversations";
 
+interface GranularNodeState {
+	nodeIdentities?: Record<string, RawNodeIdentity>;
+	nodeMemory?: Record<string, RawMemoryUsage>;
+	nodeSystem?: Record<string, RawSystemPerformanceProfile>;
+	nodeNetwork?: Record<string, RawNodeNetworkInfo>;
+}
+
+function transformNetworkInterface(iface: RawNetworkInterfaceInfo): {
+	name?: string;
+	addresses: string[];
+} {
+	const addresses: string[] = [];
+	if (iface.ipAddress && typeof iface.ipAddress === "string") {
+		addresses.push(iface.ipAddress);
+	}
+	if (Array.isArray(iface.addresses)) {
+		for (const addr of iface.addresses) {
+			if (typeof addr === "string") addresses.push(addr);
+			else if (addr && typeof addr === "object" && addr.address)
+				addresses.push(addr.address);
+		}
+	}
+	if (Array.isArray(iface.ipAddresses)) {
+		addresses.push(
+			...iface.ipAddresses.filter((a): a is string => typeof a === "string"),
+		);
+	}
+	if (Array.isArray(iface.ips)) {
+		addresses.push(
+			...iface.ips.filter((a): a is string => typeof a === "string"),
+		);
+	}
+	if (iface.ipv4 && typeof iface.ipv4 === "string") addresses.push(iface.ipv4);
+	if (iface.ipv6 && typeof iface.ipv6 === "string") addresses.push(iface.ipv6);
+
+	return {
+		name: iface.name,
+		addresses: Array.from(new Set(addresses)),
+	};
+}
+
 function transformTopology(
 	raw: RawTopology,
-	profiles?: RawNodeProfiles,
+	granularState: GranularNodeState,
 ): TopologyData {
 	const nodes: Record<string, NodeInfo> = {};
 	const edges: TopologyEdge[] = [];
 
-	for (const node of raw.nodes || []) {
-		const mergedProfile = profiles?.[node.nodeId];
-		const profile = { ...(node.nodeProfile ?? {}), ...(mergedProfile ?? {}) };
-		const ramTotal = profile?.memory?.ramTotal?.inBytes ?? 0;
-		const ramAvailable = profile?.memory?.ramAvailable?.inBytes ?? 0;
+	for (const nodeId of raw.nodes || []) {
+		if (!nodeId) continue;
+
+		// Get data from granular state mappings
+		const identity = granularState.nodeIdentities?.[nodeId];
+		const memory = granularState.nodeMemory?.[nodeId];
+		const system = granularState.nodeSystem?.[nodeId];
+		const network = granularState.nodeNetwork?.[nodeId];
+
+		const ramTotal = memory?.ramTotal?.inBytes ?? 0;
+		const ramAvailable = memory?.ramAvailable?.inBytes ?? 0;
 		const ramUsage = Math.max(ramTotal - ramAvailable, 0);
 
-		const networkInterfaces = (profile?.networkInterfaces || []).map(
-			(iface) => {
-				const addresses: string[] = [];
-				if (iface.ipAddress && typeof iface.ipAddress === "string") {
-					addresses.push(iface.ipAddress);
-				}
-				if (Array.isArray(iface.addresses)) {
-					for (const addr of iface.addresses) {
-						if (typeof addr === "string") addresses.push(addr);
-						else if (addr && typeof addr === "object" && addr.address)
-							addresses.push(addr.address);
-					}
-				}
-				if (Array.isArray(iface.ipAddresses)) {
-					addresses.push(
-						...iface.ipAddresses.filter(
-							(a): a is string => typeof a === "string",
-						),
-					);
-				}
-				if (Array.isArray(iface.ips)) {
-					addresses.push(
-						...iface.ips.filter((a): a is string => typeof a === "string"),
-					);
-				}
-				if (iface.ipv4 && typeof iface.ipv4 === "string")
-					addresses.push(iface.ipv4);
-				if (iface.ipv6 && typeof iface.ipv6 === "string")
-					addresses.push(iface.ipv6);
-
-				return {
-					name: iface.name,
-					addresses: Array.from(new Set(addresses)),
-				};
-			},
-		);
+		const rawInterfaces = network?.interfaces || [];
+		const networkInterfaces = rawInterfaces.map(transformNetworkInterface);
 
 		const ipToInterface: Record<string, string> = {};
 		for (const iface of networkInterfaces) {
@@ -264,10 +295,10 @@ function transformTopology(
 			}
 		}
 
-		nodes[node.nodeId] = {
+		nodes[nodeId] = {
 			system_info: {
-				model_id: profile?.modelId ?? "Unknown",
-				chip: profile?.chipId,
+				model_id: identity?.modelId ?? "Unknown",
+				chip: identity?.chipId,
 				memory: ramTotal,
 			},
 			network_interfaces: networkInterfaces,
@@ -278,43 +309,42 @@ function transformTopology(
 					ram_total: ramTotal,
 				},
 				temp:
-					profile?.system?.temp !== undefined
-						? { gpu_temp_avg: profile.system.temp }
+					system?.temp !== undefined
+						? { gpu_temp_avg: system.temp }
 						: undefined,
 				gpu_usage:
-					profile?.system?.gpuUsage !== undefined
-						? [0, profile.system.gpuUsage]
-						: undefined,
-				sys_power: profile?.system?.sysPower,
+					system?.gpuUsage !== undefined ? [0, system.gpuUsage] : undefined,
+				sys_power: system?.sysPower,
 			},
 			last_macmon_update: Date.now() / 1000,
-			friendly_name: profile?.friendlyName,
+			friendly_name: identity?.friendlyName,
 		};
 	}
 
-	for (const conn of raw.connections || []) {
-		if (!conn.localNodeId || !conn.sendBackNodeId) continue;
-		if (conn.localNodeId === conn.sendBackNodeId) continue;
-		if (!nodes[conn.localNodeId] || !nodes[conn.sendBackNodeId]) continue;
+	// Handle connections - nested mapping format { source: { sink: [edges] } }
+	const connections = raw.connections;
+	if (connections && typeof connections === "object") {
+		for (const [source, sinks] of Object.entries(connections)) {
+			if (!sinks || typeof sinks !== "object") continue;
+			for (const [sink, edgeList] of Object.entries(sinks)) {
+				if (!Array.isArray(edgeList)) continue;
+				for (const edge of edgeList) {
+					let sendBackIp: string | undefined;
+					if (edge && typeof edge === "object" && "sinkMultiaddr" in edge) {
+						const multiaddr = edge.sinkMultiaddr;
+						if (multiaddr) {
+							sendBackIp =
+								multiaddr.ip_address ||
+								extractIpFromMultiaddr(multiaddr.address);
+						}
+					}
 
-		let sendBackIp: string | undefined;
-		if (conn.sendBackMultiaddr) {
-			const multi = conn.sendBackMultiaddr;
-			if (typeof multi === "string") {
-				sendBackIp = extractIpFromMultiaddr(multi);
-			} else {
-				sendBackIp =
-					multi.ip_address ||
-					extractIpFromMultiaddr(multi.multiaddr) ||
-					extractIpFromMultiaddr(multi.address);
+					if (nodes[source] && nodes[sink] && source !== sink) {
+						edges.push({ source, target: sink, sendBackIp });
+					}
+				}
 			}
 		}
-
-		edges.push({
-			source: conn.localNodeId,
-			target: conn.sendBackNodeId,
-			sendBackIp,
-		});
 	}
 
 	return { nodes, edges };
@@ -868,7 +898,12 @@ class AppStore {
 			const data: RawStateResponse = await response.json();
 
 			if (data.topology) {
-				this.topologyData = transformTopology(data.topology, data.nodeProfiles);
+				this.topologyData = transformTopology(data.topology, {
+					nodeIdentities: data.nodeIdentities,
+					nodeMemory: data.nodeMemory,
+					nodeSystem: data.nodeSystem,
+					nodeNetwork: data.nodeNetwork,
+				});
 			}
 			if (data.instances) {
 				this.instances = data.instances;

dashboard/src/routes/+page.svelte

@@ -434,8 +434,8 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 		const shardData = shardObj[shardKeys[0]] as Record<string, unknown>;
 		if (!shardData) return null;
 		
-		// Model meta is nested: shard.model_meta.model_id
-		const modelMeta = shardData.model_meta ?? shardData.modelMeta;
+		// Model meta is nested: shard.model_card.model_id
+		const modelMeta = shardData.model_card ?? shardData.modelCard;
 		if (!modelMeta || typeof modelMeta !== 'object') return null;
 		
 		const meta = modelMeta as Record<string, unknown>;
@@ -915,7 +915,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 		const runnerEntries = Object.entries(runnerToShard).map(([runnerId, shardWrapped]) => {
 			const [tag, shard] = getTagged(shardWrapped);
 			const meta = (shard as { modelMeta?: { worldSize?: number; nLayers?: number; deviceRank?: number } } | undefined);
-			const deviceRank = (meta?.deviceRank as number | undefined) ?? 0;
+			const deviceRank = meta?.modelMeta?.deviceRank ?? 0;
 			return { runnerId, tag, deviceRank };
 		});
 

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

@@ -98,7 +98,7 @@
 		const shardData = shardObj[shardKeys[0]] as Record<string, unknown>;
 		if (!shardData) return null;
 
-		const modelMeta = shardData.model_meta ?? shardData.modelMeta;
+		const modelMeta = shardData.model_card ?? shardData.modelCard;
 		if (!modelMeta || typeof modelMeta !== 'object') return null;
 
 		const meta = modelMeta as Record<string, unknown>;
@@ -190,7 +190,7 @@
 						const shardKeys = Object.keys(shardObj);
 						if (shardKeys.length !== 1) return null;
 						const shardData = shardObj[shardKeys[0]] as Record<string, unknown>;
-						const modelMeta = shardData?.model_meta ?? shardData?.modelMeta;
+						const modelMeta = shardData?.model_card ?? shardData?.modelCard;
 						if (!modelMeta || typeof modelMeta !== 'object') return null;
 						const meta = modelMeta as Record<string, unknown>;
 						return (meta.prettyName as string) ?? null;

pyproject.toml

@@ -17,12 +17,13 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx==0.30.1; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.1; sys_platform == 'linux'",
+    "mlx==0.30.3; sys_platform == 'darwin'",
+    "mlx[cpu]==0.30.3; sys_platform == 'linux'",
     "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",
 ]
 
 [project.scripts]
@@ -125,3 +126,6 @@ env = [
   "EXO_TESTS=1"
 ]
 addopts = "-m 'not slow'"
+filterwarnings = [
+    "ignore:builtin type Swig:DeprecationWarning",
+]

src/exo/main.py

@@ -205,6 +205,14 @@ def main():
     logger.info("Starting EXO")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
+    # Set FAST_SYNCH override env var for runner subprocesses
+    if args.fast_synch is True:
+        os.environ["EXO_FAST_SYNCH"] = "on"
+        logger.info("FAST_SYNCH forced ON")
+    elif args.fast_synch is False:
+        os.environ["EXO_FAST_SYNCH"] = "off"
+        logger.info("FAST_SYNCH forced OFF")
+
     node = anyio.run(Node.create, args)
     anyio.run(node.run)
     logger.info("EXO Shutdown complete")
@@ -218,6 +226,7 @@ class Args(CamelCaseModel):
     api_port: PositiveInt = 52415
     tb_only: bool = False
     no_worker: bool = False
+    fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
     def parse(cls) -> Self:
@@ -259,6 +268,20 @@ class Args(CamelCaseModel):
             "--no-worker",
             action="store_true",
         )
+        fast_synch_group = parser.add_mutually_exclusive_group()
+        fast_synch_group.add_argument(
+            "--fast-synch",
+            action="store_true",
+            dest="fast_synch",
+            default=None,
+            help="Force MLX FAST_SYNCH on (for JACCL backend)",
+        )
+        fast_synch_group.add_argument(
+            "--no-fast-synch",
+            action="store_false",
+            dest="fast_synch",
+            help="Force MLX FAST_SYNCH off",
+        )
 
         args = parser.parse_args()
         return cls(**vars(args))  # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically

src/exo/master/api.py

@@ -1,13 +1,14 @@
 import time
 from collections.abc import AsyncGenerator
+from http import HTTPStatus
 from typing import cast
 
 import anyio
-from anyio import create_task_group
+from anyio import BrokenResourceError, create_task_group
 from anyio.abc import TaskGroup
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType]
 from hypercorn.config import Config
@@ -18,8 +19,8 @@ from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
-from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.models.model_meta import get_model_meta
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
+from exo.shared.models.model_meta import get_model_card
 from exo.shared.types.api import (
     BenchChatCompletionResponse,
     BenchChatCompletionTaskParams,
@@ -29,6 +30,8 @@ from exo.shared.types.api import (
     CreateInstanceParams,
     CreateInstanceResponse,
     DeleteInstanceResponse,
+    ErrorInfo,
+    ErrorResponse,
     FinishReason,
     GenerationStats,
     ModelList,
@@ -49,9 +52,13 @@ from exo.shared.types.commands import (
     TaskFinished,
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId
-from exo.shared.types.events import ChunkGenerated, Event, ForwarderEvent, IndexedEvent
+from exo.shared.types.events import (
+    ChunkGenerated,
+    Event,
+    ForwarderEvent,
+    IndexedEvent,
+)
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.state import State
 from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -79,12 +86,12 @@ def chunk_to_response(
     )
 
 
-async def resolve_model_meta(model_id: str) -> ModelMetadata:
+async def resolve_model_card(model_id: str) -> ModelCard:
     if model_id in MODEL_CARDS:
         model_card = MODEL_CARDS[model_id]
-        return model_card.metadata
+        return model_card
     else:
-        return await get_model_meta(model_id)
+        return await get_model_card(model_id)
 
 
 class API:
@@ -115,6 +122,7 @@ class API:
         self.paused_ev: anyio.Event = anyio.Event()
 
         self.app = FastAPI()
+        self._setup_exception_handlers()
         self._setup_cors()
         self._setup_routes()
 
@@ -145,6 +153,21 @@ class API:
         self.paused_ev.set()
         self.paused_ev = anyio.Event()
 
+    def _setup_exception_handlers(self) -> None:
+        self.app.exception_handler(HTTPException)(self.http_exception_handler)
+
+    async def http_exception_handler(
+        self, _: Request, exc: HTTPException
+    ) -> JSONResponse:
+        err = ErrorResponse(
+            error=ErrorInfo(
+                message=exc.detail,
+                type=HTTPStatus(exc.status_code).phrase,
+                code=exc.status_code,
+            )
+        )
+        return JSONResponse(err.model_dump(), status_code=exc.status_code)
+
     def _setup_cors(self) -> None:
         self.app.add_middleware(
             CORSMiddleware,
@@ -173,7 +196,7 @@ class API:
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
-            model_meta=await resolve_model_meta(payload.model_id),
+            model_card=await resolve_model_card(payload.model_id),
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
@@ -183,15 +206,15 @@ class API:
         return CreateInstanceResponse(
             message="Command received.",
             command_id=command.command_id,
-            model_meta=command.model_meta,
+            model_card=command.model_card,
         )
 
     async def create_instance(
         self, payload: CreateInstanceParams
     ) -> CreateInstanceResponse:
         instance = payload.instance
-        model_meta = await resolve_model_meta(instance.shard_assignments.model_id)
-        required_memory = model_meta.storage_size
+        model_card = await resolve_model_card(instance.shard_assignments.model_id)
+        required_memory = model_card.storage_size
         available_memory = self._calculate_total_available_memory()
 
         if required_memory > available_memory:
@@ -208,7 +231,7 @@ class API:
         return CreateInstanceResponse(
             message="Command received.",
             command_id=command.command_id,
-            model_meta=model_meta,
+            model_card=model_card,
         )
 
     async def get_placement(
@@ -218,16 +241,18 @@ class API:
         instance_meta: InstanceMeta = InstanceMeta.MlxRing,
         min_nodes: int = 1,
     ) -> Instance:
-        model_meta = await resolve_model_meta(model_id)
+        model_card = await resolve_model_card(model_id)
 
         try:
             placements = get_instance_placements(
                 PlaceInstance(
-                    model_meta=model_meta,
+                    model_card=model_card,
                     sharding=sharding,
                     instance_meta=instance_meta,
                     min_nodes=min_nodes,
                 ),
+                node_memory=self.state.node_memory,
+                node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
             )
@@ -254,7 +279,7 @@ class API:
         if len(list(self.state.topology.list_nodes())) == 0:
             return PlacementPreviewResponse(previews=[])
 
-        cards = [card for card in MODEL_CARDS.values() if card.short_id == model_id]
+        cards = [card for card in MODEL_CARDS.values() if card.model_id == model_id]
         if not cards:
             raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
 
@@ -272,32 +297,33 @@ class API:
         # TODO: PDD
         # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
 
-        for card in cards:
-            model_meta = card.metadata
+        for model_card in cards:
             for sharding, instance_meta, min_nodes in instance_combinations:
                 try:
                     placements = get_instance_placements(
                         PlaceInstance(
-                            model_meta=model_meta,
+                            model_card=model_card,
                             sharding=sharding,
                             instance_meta=instance_meta,
                             min_nodes=min_nodes,
                         ),
+                        node_memory=self.state.node_memory,
+                        node_network=self.state.node_network,
                         topology=self.state.topology,
                         current_instances=self.state.instances,
                     )
                 except ValueError as exc:
-                    if (card.model_id, sharding, instance_meta, 0) not in seen:
+                    if (model_card.model_id, sharding, instance_meta, 0) not in seen:
                         previews.append(
                             PlacementPreview(
-                                model_id=card.model_id,
+                                model_id=model_card.model_id,
                                 sharding=sharding,
                                 instance_meta=instance_meta,
                                 instance=None,
                                 error=str(exc),
                             )
                         )
-                    seen.add((card.model_id, sharding, instance_meta, 0))
+                    seen.add((model_card.model_id, sharding, instance_meta, 0))
                     continue
 
                 current_ids = set(self.state.instances.keys())
@@ -308,17 +334,17 @@ class API:
                 ]
 
                 if len(new_instances) != 1:
-                    if (card.model_id, sharding, instance_meta, 0) not in seen:
+                    if (model_card.model_id, sharding, instance_meta, 0) not in seen:
                         previews.append(
                             PlacementPreview(
-                                model_id=card.model_id,
+                                model_id=model_card.model_id,
                                 sharding=sharding,
                                 instance_meta=instance_meta,
                                 instance=None,
                                 error="Expected exactly one new instance from placement",
                             )
                         )
-                    seen.add((card.model_id, sharding, instance_meta, 0))
+                    seen.add((model_card.model_id, sharding, instance_meta, 0))
                     continue
 
                 instance = new_instances[0]
@@ -327,7 +353,7 @@ class API:
 
                 memory_delta_by_node: dict[str, int] = {}
                 if node_ids:
-                    total_bytes = model_meta.storage_size.in_bytes
+                    total_bytes = model_card.storage_size.in_bytes
                     per_node = total_bytes // len(node_ids)
                     remainder = total_bytes % len(node_ids)
                     for index, node_id in enumerate(sorted(node_ids, key=str)):
@@ -335,14 +361,14 @@ class API:
                         memory_delta_by_node[str(node_id)] = per_node + extra
 
                 if (
-                    card.model_id,
+                    model_card.model_id,
                     sharding,
                     instance_meta,
                     len(node_ids),
                 ) not in seen:
                     previews.append(
                         PlacementPreview(
-                            model_id=card.model_id,
+                            model_id=model_card.model_id,
                             sharding=sharding,
                             instance_meta=instance_meta,
                             instance=instance,
@@ -350,7 +376,7 @@ class API:
                             error=None,
                         )
                     )
-                seen.add((card.model_id, sharding, instance_meta, len(node_ids)))
+                seen.add((model_card.model_id, sharding, instance_meta, len(node_ids)))
 
         return PlacementPreviewResponse(previews=previews)
 
@@ -406,6 +432,18 @@ class API:
         """Generate chat completion stream as JSON strings."""
 
         async for chunk in self._chat_chunk_stream(command_id):
+            if chunk.finish_reason == "error":
+                error_response = ErrorResponse(
+                    error=ErrorInfo(
+                        message=chunk.error_message or "Internal server error",
+                        type="InternalServerError",
+                        code=500,
+                    )
+                )
+                yield f"data: {error_response.model_dump_json()}\n\n"
+                yield "data: [DONE]\n\n"
+                return
+
             chunk_response: ChatCompletionResponse = chunk_to_response(
                 chunk, command_id
             )
@@ -426,6 +464,12 @@ class API:
         finish_reason: FinishReason | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
+            if chunk.finish_reason == "error":
+                raise HTTPException(
+                    status_code=500,
+                    detail=chunk.error_message or "Internal server error",
+                )
+
             if model is None:
                 model = chunk.model
 
@@ -463,6 +507,12 @@ class API:
         stats: GenerationStats | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
+            if chunk.finish_reason == "error":
+                raise HTTPException(
+                    status_code=500,
+                    detail=chunk.error_message or "Internal server error",
+                )
+
             if model is None:
                 model = chunk.model
 
@@ -501,8 +551,8 @@ class API:
         self, payload: ChatCompletionTaskParams
     ) -> ChatCompletionResponse | StreamingResponse:
         """Handle chat completions, supporting both streaming and non-streaming responses."""
-        model_meta = await resolve_model_meta(payload.model)
-        payload.model = model_meta.model_id
+        model_card = await resolve_model_card(payload.model)
+        payload.model = model_card.model_id
 
         if not any(
             instance.shard_assignments.model_id == payload.model
@@ -528,8 +578,8 @@ class API:
     async def bench_chat_completions(
         self, payload: BenchChatCompletionTaskParams
     ) -> BenchChatCompletionResponse:
-        model_meta = await resolve_model_meta(payload.model)
-        payload.model = model_meta.model_id
+        model_card = await resolve_model_card(payload.model)
+        payload.model = model_card.model_id
 
         if not any(
             instance.shard_assignments.model_id == payload.model
@@ -552,9 +602,8 @@ class API:
         """Calculate total available memory across all nodes in bytes."""
         total_available = Memory()
 
-        for node in self.state.topology.list_nodes():
-            if node.node_profile is not None:
-                total_available += node.node_profile.memory.ram_available
+        for memory in self.state.node_memory.values():
+            total_available += memory.ram_available
 
         return total_available
 
@@ -563,13 +612,13 @@ class API:
         return ModelList(
             data=[
                 ModelListModel(
-                    id=card.short_id,
+                    id=card.model_id,
                     hugging_face_id=card.model_id,
-                    name=card.name,
-                    description=card.description,
-                    tags=card.tags,
-                    storage_size_megabytes=int(card.metadata.storage_size.in_mb),
-                    supports_tensor=card.metadata.supports_tensor,
+                    name=card.model_id.short(),
+                    description="",
+                    tags=[],
+                    storage_size_megabytes=int(card.storage_size.in_mb),
+                    supports_tensor=card.supports_tensor,
                 )
                 for card in MODEL_CARDS.values()
             ]
@@ -607,14 +656,14 @@ class API:
                 for idx, event in self.event_buffer.drain_indexed():
                     self._event_log.append(event)
                     self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
-                    if (
-                        isinstance(event, ChunkGenerated)
-                        and event.command_id in self._chat_completion_queues
-                    ):
+                    if isinstance(event, ChunkGenerated):
                         assert isinstance(event.chunk, TokenChunk)
-                        await self._chat_completion_queues[event.command_id].send(
-                            event.chunk
-                        )
+                        queue = self._chat_completion_queues.get(event.command_id)
+                        if queue is not None:
+                            try:
+                                await queue.send(event.chunk)
+                            except BrokenResourceError:
+                                self._chat_completion_queues.pop(event.command_id, None)
 
     async def _pause_on_new_election(self):
         with self.election_receiver as ems:

src/exo/master/main.py

@@ -27,6 +27,7 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InstanceDeleted,
+    NodeGatheredInfo,
     NodeTimedOut,
     TaskCreated,
     TaskDeleted,
@@ -158,6 +159,8 @@ class Master:
                                 command,
                                 self.state.topology,
                                 self.state.instances,
+                                self.state.node_memory,
+                                self.state.node_network,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement
@@ -200,9 +203,7 @@ class Master:
     async def _plan(self) -> None:
         while True:
             # kill broken instances
-            connected_node_ids = set(
-                [x.node_id for x in self.state.topology.list_nodes()]
-            )
+            connected_node_ids = set(self.state.topology.list_nodes())
             for instance_id, instance in self.state.instances.items():
                 for node_id in instance.shard_assignments.node_to_runner:
                     if node_id not in connected_node_ids:
@@ -237,6 +238,8 @@ class Master:
                     self.state = apply(self.state, indexed)
 
                     event._master_time_stamp = datetime.now(tz=timezone.utc)  # pyright: ignore[reportPrivateUsage]
+                    if isinstance(event, NodeGatheredInfo):
+                        event.when = str(datetime.now(tz=timezone.utc))
 
                     self._event_log.append(event)
                     await self._send_event(indexed)

src/exo/master/placement.py

@@ -6,23 +6,25 @@ from typing import Sequence
 from loguru import logger
 
 from exo.master.placement_utils import (
+    Cycle,
     filter_cycles_by_memory,
-    get_mlx_ibv_devices_matrix,
     get_mlx_jaccl_coordinators,
+    get_mlx_jaccl_devices_matrix,
     get_mlx_ring_hosts_by_node,
     get_shard_assignments,
     get_smallest_cycles,
 )
+from exo.shared.models.model_cards import ModelId
 from exo.shared.topology import Topology
 from exo.shared.types.commands import (
     CreateInstance,
     DeleteInstance,
     PlaceInstance,
 )
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId
-from exo.shared.types.topology import NodeInfo
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.worker.instances import (
     Instance,
     InstanceId,
@@ -52,37 +54,33 @@ def place_instance(
     command: PlaceInstance,
     topology: Topology,
     current_instances: Mapping[InstanceId, Instance],
+    node_memory: Mapping[NodeId, MemoryUsage],
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[InstanceId, Instance]:
-    all_nodes = list(topology.list_nodes())
-
-    logger.info("finding cycles:")
     cycles = topology.get_cycles()
-    singleton_cycles = [[node] for node in all_nodes]
-    candidate_cycles = list(
-        filter(lambda it: len(it) >= command.min_nodes, cycles + singleton_cycles)
-    )
+    candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
     cycles_with_sufficient_memory = filter_cycles_by_memory(
-        candidate_cycles, command.model_meta.storage_size
+        candidate_cycles, node_memory, command.model_card.storage_size
     )
-    if not cycles_with_sufficient_memory:
+    if len(cycles_with_sufficient_memory) == 0:
         raise ValueError("No cycles found with sufficient memory")
 
     if command.sharding == Sharding.Tensor:
-        if not command.model_meta.supports_tensor:
+        if not command.model_card.supports_tensor:
             raise ValueError(
-                f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_meta.model_id}"
+                f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}"
             )
         # TODO: the condition here for tensor parallel is not correct, but it works good enough for now.
         cycles_with_sufficient_memory = [
             cycle
             for cycle in cycles_with_sufficient_memory
-            if command.model_meta.hidden_size % len(cycle) == 0
+            if command.model_card.hidden_size % len(cycle) == 0
         ]
         if not cycles_with_sufficient_memory:
             raise ValueError(
-                f"No tensor sharding found for model with hidden_size {command.model_meta.hidden_size} candidate cycles"
+                f"No tensor sharding found for model with hidden_size {command.model_card.hidden_size} candidate cycles"
             )
-    if command.sharding == Sharding.Pipeline and command.model_meta.model_id == ModelId(
+    if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId(
         "mlx-community/DeepSeek-V3.1-8bit"
     ):
         raise ValueError(
@@ -92,44 +90,38 @@ def place_instance(
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
     smallest_tb_cycles = [
-        cycle
-        for cycle in smallest_cycles
-        if topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
+        cycle for cycle in smallest_cycles if topology.is_thunderbolt_cycle(cycle)
     ]
 
     if smallest_tb_cycles != []:
         smallest_cycles = smallest_tb_cycles
 
-    cycles_with_leaf_nodes: list[list[NodeInfo]] = [
+    cycles_with_leaf_nodes: list[Cycle] = [
         cycle
         for cycle in smallest_cycles
-        if any(topology.node_is_leaf(node.node_id) for node in cycle)
+        if any(topology.node_is_leaf(node_id) for node_id in cycle)
     ]
 
     selected_cycle = max(
         cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
         key=lambda cycle: sum(
-            (
-                node.node_profile.memory.ram_available
-                for node in cycle
-                if node.node_profile is not None
-            ),
+            (node_memory[node_id].ram_available for node_id in cycle),
             start=Memory(),
         ),
     )
 
     shard_assignments = get_shard_assignments(
-        command.model_meta, selected_cycle, command.sharding
+        command.model_card, selected_cycle, command.sharding, node_memory
     )
 
-    cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle)
+    cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)
 
     instance_id = InstanceId()
     target_instances = dict(deepcopy(current_instances))
 
     if len(selected_cycle) == 1:
         logger.warning(
-            "You have likely selected ibv for a single node instance; falling back to MlxRing"
+            "You have likely selected jaccl for a single node instance; falling back to MlxRing"
         )
 
         command.instance_meta = InstanceMeta.MlxRing
@@ -137,19 +129,20 @@ def place_instance(
     # TODO: Single node instances
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
-            mlx_ibv_devices = get_mlx_ibv_devices_matrix(
-                selected_cycle,
+            mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
+                [node_id for node_id in selected_cycle],
                 cycle_digraph,
             )
             mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
-                selected_cycle,
+                coordinator=selected_cycle.node_ids[0],
                 coordinator_port=random_ephemeral_port(),
                 cycle_digraph=cycle_digraph,
+                node_network=node_network,
             )
             target_instances[instance_id] = MlxJacclInstance(
                 instance_id=instance_id,
                 shard_assignments=shard_assignments,
-                ibv_devices=mlx_ibv_devices,
+                jaccl_devices=mlx_jaccl_devices,
                 jaccl_coordinators=mlx_jaccl_coordinators,
             )
         case InstanceMeta.MlxRing:
@@ -158,6 +151,7 @@ def place_instance(
                 selected_cycle=selected_cycle,
                 cycle_digraph=cycle_digraph,
                 ephemeral_port=ephemeral_port,
+                node_network=node_network,
             )
             target_instances[instance_id] = MlxRingInstance(
                 instance_id=instance_id,

src/exo/master/placement_utils.py

@@ -1,15 +1,13 @@
-from collections.abc import Generator
-from typing import TypeGuard, cast
+from collections.abc import Generator, Mapping
 
 from loguru import logger
-from pydantic import BaseModel
 
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelMetadata
-from exo.shared.types.profiling import NodePerformanceProfile
-from exo.shared.types.topology import NodeInfo
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
@@ -19,67 +17,113 @@ from exo.shared.types.worker.shards import (
 )
 
 
-class NodeWithProfile(BaseModel):
-    node_id: NodeId
-    node_profile: NodePerformanceProfile
-
-
-def narrow_all_nodes(nodes: list[NodeInfo]) -> TypeGuard[list[NodeWithProfile]]:
-    return all(node.node_profile is not None for node in nodes)
-
-
 def filter_cycles_by_memory(
-    cycles: list[list[NodeInfo]], required_memory: Memory
-) -> list[list[NodeInfo]]:
-    filtered_cycles: list[list[NodeInfo]] = []
+    cycles: list[Cycle],
+    node_memory: Mapping[NodeId, MemoryUsage],
+    required_memory: Memory,
+) -> list[Cycle]:
+    filtered_cycles: list[Cycle] = []
     for cycle in cycles:
-        if not narrow_all_nodes(cycle):
+        if not all(node in node_memory for node in cycle):
             continue
 
         total_mem = sum(
-            (node.node_profile.memory.ram_available for node in cycle), start=Memory()
+            (node_memory[node_id].ram_available for node_id in cycle.node_ids),
+            start=Memory(),
         )
         if total_mem >= required_memory:
-            filtered_cycles.append(cast(list[NodeInfo], cycle))
+            filtered_cycles.append(cycle)
     return filtered_cycles
 
 
-def get_smallest_cycles(cycles: list[list[NodeInfo]]) -> list[list[NodeInfo]]:
+def get_smallest_cycles(
+    cycles: list[Cycle],
+) -> list[Cycle]:
     min_nodes = min(len(cycle) for cycle in cycles)
     return [cycle for cycle in cycles if len(cycle) == min_nodes]
 
 
+def allocate_layers_proportionally(
+    total_layers: int,
+    memory_fractions: list[float],
+) -> list[int]:
+    n = len(memory_fractions)
+    if n == 0:
+        raise ValueError("Cannot allocate layers to an empty node list")
+    if total_layers < n:
+        raise ValueError(
+            f"Cannot distribute {total_layers} layers across {n} nodes "
+            "(need at least 1 layer per node)"
+        )
+
+    # Largest remainder: floor each, then distribute remainder by fractional part
+    raw = [f * total_layers for f in memory_fractions]
+    result = [int(r) for r in raw]
+    by_remainder = sorted(range(n), key=lambda i: raw[i] - result[i], reverse=True)
+    for i in range(total_layers - sum(result)):
+        result[by_remainder[i]] += 1
+
+    # Ensure minimum 1 per node by taking from the largest
+    for i in range(n):
+        if result[i] == 0:
+            max_idx = max(range(n), key=lambda j: result[j])
+            assert result[max_idx] > 1
+            result[max_idx] -= 1
+            result[i] = 1
+
+    return result
+
+
 def get_shard_assignments_for_pipeline_parallel(
-    model_meta: ModelMetadata,
-    selected_cycle: list[NodeWithProfile],
+    model_card: ModelCard,
+    cycle: Cycle,
+    node_memory: Mapping[NodeId, MemoryUsage],
 ):
+    if not cycle.node_ids:
+        raise ValueError("Cannot create shard assignments for empty node cycle")
+
     cycle_memory = sum(
-        (node.node_profile.memory.ram_available for node in selected_cycle),
+        (node_memory[node_id].ram_available for node_id in cycle.node_ids),
         start=Memory(),
     )
-    total_layers = model_meta.n_layers
-    world_size = len(selected_cycle)
+    if cycle_memory.in_bytes == 0:
+        raise ValueError("Cannot create shard assignments: total available memory is 0")
+
+    total_layers = model_card.n_layers
+    world_size = len(cycle)
     runner_to_shard: dict[RunnerId, ShardMetadata] = {}
     node_to_runner: dict[NodeId, RunnerId] = {}
 
-    layers_assigned = 0
-    for i, node in enumerate(selected_cycle):
-        if i == len(selected_cycle) - 1:
-            node_layers = total_layers - layers_assigned
-        else:
-            node_layers = round(
-                total_layers
-                * (
-                    node.node_profile.memory.ram_available.in_bytes
-                    / cycle_memory.in_bytes
-                )
-            )
-            node_layers = max(1, node_layers)
+    layer_allocations = allocate_layers_proportionally(
+        total_layers=total_layers,
+        memory_fractions=[
+            node_memory[node_id].ram_available.in_bytes / cycle_memory.in_bytes
+            for node_id in cycle.node_ids
+        ],
+    )
 
+    # Validate each node has sufficient memory for its assigned layers
+    memory_per_layer = model_card.storage_size.in_bytes / total_layers
+    for i, (node_id, node_layers) in enumerate(
+        zip(cycle.node_ids, layer_allocations, strict=True)
+    ):
+        required_memory = node_layers * memory_per_layer
+        available_memory = node_memory[node_id].ram_available.in_bytes
+        if required_memory > available_memory:
+            raise ValueError(
+                f"Node {i} ({node_id}) has insufficient memory: "
+                f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
+                f"but only has {available_memory / (1024**3):.2f} GB available"
+            )
+
+    layers_assigned = 0
+    for i, (node_id, node_layers) in enumerate(
+        zip(cycle.node_ids, layer_allocations, strict=True)
+    ):
         runner_id = RunnerId()
 
         shard = PipelineShardMetadata(
-            model_meta=model_meta,
+            model_card=model_card,
             device_rank=i,
             world_size=world_size,
             start_layer=layers_assigned,
@@ -88,11 +132,11 @@ def get_shard_assignments_for_pipeline_parallel(
         )
 
         runner_to_shard[runner_id] = shard
-        node_to_runner[node.node_id] = runner_id
+        node_to_runner[node_id] = runner_id
         layers_assigned += node_layers
 
     shard_assignments = ShardAssignments(
-        model_id=model_meta.model_id,
+        model_id=model_card.model_id,
         runner_to_shard=runner_to_shard,
         node_to_runner=node_to_runner,
     )
@@ -101,17 +145,17 @@ def get_shard_assignments_for_pipeline_parallel(
 
 
 def get_shard_assignments_for_tensor_parallel(
-    model_meta: ModelMetadata,
-    selected_cycle: list[NodeWithProfile],
+    model_card: ModelCard,
+    cycle: Cycle,
 ):
-    total_layers = model_meta.n_layers
-    world_size = len(selected_cycle)
+    total_layers = model_card.n_layers
+    world_size = len(cycle)
     runner_to_shard: dict[RunnerId, ShardMetadata] = {}
     node_to_runner: dict[NodeId, RunnerId] = {}
 
-    for i, node in enumerate(selected_cycle):
+    for i, node_id in enumerate(cycle):
         shard = TensorShardMetadata(
-            model_meta=model_meta,
+            model_card=model_card,
             device_rank=i,
             world_size=world_size,
             start_layer=0,
@@ -122,10 +166,10 @@ def get_shard_assignments_for_tensor_parallel(
         runner_id = RunnerId()
 
         runner_to_shard[runner_id] = shard
-        node_to_runner[node.node_id] = runner_id
+        node_to_runner[node_id] = runner_id
 
     shard_assignments = ShardAssignments(
-        model_id=model_meta.model_id,
+        model_id=model_card.model_id,
         runner_to_shard=runner_to_shard,
         node_to_runner=node_to_runner,
     )
@@ -134,22 +178,22 @@ def get_shard_assignments_for_tensor_parallel(
 
 
 def get_shard_assignments(
-    model_meta: ModelMetadata,
-    selected_cycle: list[NodeInfo],
+    model_card: ModelCard,
+    cycle: Cycle,
     sharding: Sharding,
+    node_memory: Mapping[NodeId, MemoryUsage],
 ) -> ShardAssignments:
-    if not narrow_all_nodes(selected_cycle):
-        raise ValueError("All nodes must have profiles to create shard assignments")
     match sharding:
         case Sharding.Pipeline:
             return get_shard_assignments_for_pipeline_parallel(
-                model_meta=model_meta,
-                selected_cycle=selected_cycle,
+                model_card=model_card,
+                cycle=cycle,
+                node_memory=node_memory,
             )
         case Sharding.Tensor:
             return get_shard_assignments_for_tensor_parallel(
-                model_meta=model_meta,
-                selected_cycle=selected_cycle,
+                model_card=model_card,
+                cycle=cycle,
             )
 
 
@@ -164,38 +208,40 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
             )
         return []
 
+    cycle = cycles[0]
+
     get_thunderbolt = False
-    if cycle_digraph.is_thunderbolt_cycle(cycles[0]):
+    if cycle_digraph.is_thunderbolt_cycle(cycle):
         get_thunderbolt = True
 
     logger.info(f"Using thunderbolt cycle: {get_thunderbolt}")
 
-    cycle = cycles[0]
     hosts: list[Host] = []
     for i in range(len(cycle)):
-        current_node = cycle[i]
-        next_node = cycle[(i + 1) % len(cycle)]
+        current_node = cycle.node_ids[i]
+        next_node = cycle.node_ids[(i + 1) % len(cycle)]
 
-        for connection in cycle_digraph.list_connections():
-            if (
-                connection.local_node_id == current_node.node_id
-                and connection.send_back_node_id == next_node.node_id
-            ):
-                if get_thunderbolt and not connection.is_thunderbolt():
-                    continue
-                assert connection.send_back_multiaddr is not None
-                host = Host(
-                    ip=connection.send_back_multiaddr.ip_address,
-                    port=connection.send_back_multiaddr.port,
-                )
-                hosts.append(host)
-                break
+        for connection in cycle_digraph.get_all_connections_between(
+            source=current_node, sink=next_node
+        ):
+            if not isinstance(connection, SocketConnection):
+                continue
+
+            if get_thunderbolt and not connection.is_thunderbolt():
+                continue
+
+            host = Host(
+                ip=connection.sink_multiaddr.ip_address,
+                port=connection.sink_multiaddr.port,
+            )
+            hosts.append(host)
+            break
 
     return hosts
 
 
-def get_mlx_ibv_devices_matrix(
-    selected_cycle: list[NodeInfo],
+def get_mlx_jaccl_devices_matrix(
+    selected_cycle: list[NodeId],
     cycle_digraph: Topology,
 ) -> list[list[str | None]]:
     """Build connectivity matrix mapping device i to device j via RDMA interface names.
@@ -214,72 +260,37 @@ def get_mlx_ibv_devices_matrix(
             if i == j:
                 continue
 
-            # Find the IP J uses to talk to I
-            for connection_ip, _ in _find_connection_ip(node_j, node_i, cycle_digraph):
-                # This is a local IP on I, which is attached to an interface: find that interface
-                if interface_name := _find_rdma_interface_name_for_ip(
-                    connection_ip, node_i
-                ):
-                    matrix[i][j] = interface_name
-                    logger.info(
-                        f"Interface name for {connection_ip} on {node_i.node_id}: {interface_name}"
-                    )
+            for conn in cycle_digraph.get_all_connections_between(node_i, node_j):
+                if isinstance(conn, RDMAConnection):
+                    matrix[i][j] = conn.source_rdma_iface
                     break
             else:
                 logger.warning(
-                    f"Failed to find interface name between {node_i.node_id} and {node_j.node_id}"
+                    f"Failed to find interface name between {node_i} and {node_j}"
                 )
                 raise ValueError(
-                    "Current ibv backend requires all-to-all rdma connections"
+                    "Current jaccl backend requires all-to-all RDMA connections"
                 )
 
     return matrix
 
 
 def _find_connection_ip(
-    node_i: NodeInfo,
-    node_j: NodeInfo,
+    node_i: NodeId,
+    node_j: NodeId,
     cycle_digraph: Topology,
 ) -> Generator[tuple[str, bool]]:
-    """Find all IP addresses that connect node i to node j, with thunderbolt flag."""
-    for connection in cycle_digraph.list_connections():
-        if (
-            connection.local_node_id == node_i.node_id
-            and connection.send_back_node_id == node_j.node_id
-        ):
-            yield connection.send_back_multiaddr.ip_address, connection.is_thunderbolt()
-
-
-def _find_rdma_interface_name_for_ip(
-    ip_address: str,
-    node_info: NodeInfo,
-) -> str | None:
-    if node_info.node_profile is None:
-        return None
-
-    logger.info(f"Searching {node_info.node_id} for ip {ip_address}:")
-    for interface in node_info.node_profile.network_interfaces:
-        if interface.name not in ["en2", "en3", "en4", "en5", "en6", "en7"]:
-            continue
-        logger.info(f" | {interface.name}: {interface.ip_address}")
-        if interface.ip_address != ip_address:
-            continue
-
-        logger.info("Found")
-        return f"rdma_{interface.name}"
-
-    return None
+    """Find all IP addresses that connect node i to node j."""
+    for connection in cycle_digraph.get_all_connections_between(node_i, node_j):
+        if isinstance(connection, SocketConnection):
+            yield connection.sink_multiaddr.ip_address, connection.is_thunderbolt()
 
 
 def _find_interface_name_for_ip(
-    ip_address: str,
-    node_info: NodeInfo,
+    ip_address: str, node_network: NodeNetworkInfo
 ) -> str | None:
     """Find the interface name for an IP address on a node (any interface)."""
-    if node_info.node_profile is None:
-        return None
-
-    for interface in node_info.node_profile.network_interfaces:
+    for interface in node_network.interfaces:
         if interface.ip_address == ip_address:
             return interface.name
 
@@ -287,7 +298,10 @@ def _find_interface_name_for_ip(
 
 
 def _find_ip_prioritised(
-    node: NodeInfo, other_node: NodeInfo, cycle_digraph: Topology
+    node_id: NodeId,
+    other_node_id: NodeId,
+    cycle_digraph: Topology,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> str | None:
     # TODO: Actually prioritize in the correct Ethernet > Wifi > Non-TB > TB order.
     """Find an IP address between nodes with prioritization.
@@ -298,9 +312,14 @@ def _find_ip_prioritised(
     3. Non-Thunderbolt connections
     4. Any other IP address
     """
-    ips = list(_find_connection_ip(node, other_node, cycle_digraph))
+    ips = list(_find_connection_ip(node_id, other_node_id, cycle_digraph))
     # We expect a unique iface -> ip mapping
-    iface_map = {_find_interface_name_for_ip(ip, other_node): ip for ip, _ in ips}
+    iface_map = {
+        _find_interface_name_for_ip(
+            ip, node_network.get(other_node_id, NodeNetworkInfo())
+        ): ip
+        for ip, _ in ips
+    }
 
     en0_ip = iface_map.get("en0")
     if en0_ip:
@@ -324,9 +343,10 @@ def _find_ip_prioritised(
 
 
 def get_mlx_ring_hosts_by_node(
-    selected_cycle: list[NodeInfo],
+    selected_cycle: Cycle,
     cycle_digraph: Topology,
     ephemeral_port: int,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[NodeId, list[Host]]:
     """Generate per-node host lists for MLX ring backend.
 
@@ -341,14 +361,13 @@ def get_mlx_ring_hosts_by_node(
 
     hosts_by_node: dict[NodeId, list[Host]] = {}
 
-    for rank, node in enumerate(selected_cycle):
-        node_id = node.node_id
+    for rank, node_id in enumerate(selected_cycle):
         left_rank = (rank - 1) % world_size
         right_rank = (rank + 1) % world_size
 
         hosts_for_node: list[Host] = []
 
-        for idx, other_node in enumerate(selected_cycle):
+        for idx, other_node_id in enumerate(selected_cycle):
             if idx == rank:
                 hosts_for_node.append(Host(ip="0.0.0.0", port=ephemeral_port))
                 continue
@@ -358,10 +377,12 @@ def get_mlx_ring_hosts_by_node(
                 hosts_for_node.append(Host(ip="198.51.100.1", port=0))
                 continue
 
-            connection_ip = _find_ip_prioritised(node, other_node, cycle_digraph)
+            connection_ip = _find_ip_prioritised(
+                node_id, other_node_id, cycle_digraph, node_network
+            )
             if connection_ip is None:
                 logger.warning(
-                    f"Failed to find prioritised connection IP between {node_id} and {other_node.node_id}"
+                    f"Failed to find prioritised connection IP between {node_id} and {other_node_id}"
                 )
                 raise ValueError(
                     "MLX ring backend requires connectivity between neighbouring nodes"
@@ -375,31 +396,34 @@ def get_mlx_ring_hosts_by_node(
 
 
 def get_mlx_jaccl_coordinators(
-    selected_cycle: list[NodeInfo],
+    coordinator: NodeId,
     coordinator_port: int,
     cycle_digraph: Topology,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[NodeId, str]:
-    """Get the coordinator addresses for MLX Jaccl (rank 0 device).
+    """Get the coordinator addresses for MLX JACCL (rank 0 device).
 
     Select an IP address that each node can reach for the rank 0 node. Returns
     address in format "X.X.X.X:PORT" per node.
     """
-    rank_0_node = selected_cycle[0]
-    logger.debug(f"Selecting coordinator from rank 0 node: {rank_0_node.node_id}")
+    logger.info(f"Selecting coordinator: {coordinator}")
 
-    def get_ip_for_node(n: NodeInfo) -> str:
-        if n.node_id == rank_0_node.node_id:
+    def get_ip_for_node(n: NodeId) -> str:
+        if n == coordinator:
             return "0.0.0.0"
 
-        ip = _find_ip_prioritised(n, rank_0_node, cycle_digraph)
-        if ip:
+        ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network)
+        if ip is not None:
             return ip
 
         logger.warning(
-            f"Failed to find directly connected ip between {n.node_id} and {rank_0_node.node_id}"
+            f"Failed to find directly connected ip between {n} and {coordinator}"
+        )
+        raise ValueError(
+            "Current jaccl backend requires all participating devices to be able to communicate"
         )
-        raise ValueError("Current ibv backend requires all-to-all rdma connections")
 
     return {
-        n.node_id: f"{get_ip_for_node(n)}:{coordinator_port}" for n in selected_cycle
+        n: f"{get_ip_for_node(n)}:{coordinator_port}"
+        for n in cycle_digraph.list_nodes()
     }

src/exo/master/tests/conftest.py

@@ -1,67 +1,37 @@
-from typing import Callable
-
-import pytest
-
-from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
+    MemoryUsage,
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
 )
-from exo.shared.types.topology import Connection, ConnectionProfile, NodeInfo
+from exo.shared.types.topology import RDMAConnection, SocketConnection
 
 
-@pytest.fixture
-def create_node():
-    def _create_node(memory: int, node_id: NodeId | None = None) -> NodeInfo:
-        if node_id is None:
-            node_id = NodeId()
-        return NodeInfo(
-            node_id=node_id,
-            node_profile=NodePerformanceProfile(
-                model_id="test",
-                chip_id="test",
-                friendly_name="test",
-                memory=MemoryPerformanceProfile.from_bytes(
-                    ram_total=1000,
-                    ram_available=memory,
-                    swap_total=1000,
-                    swap_available=1000,
-                ),
-                network_interfaces=[],
-                system=SystemPerformanceProfile(),
-            ),
-        )
-
-    return _create_node
+def create_node_memory(memory: int) -> MemoryUsage:
+    return MemoryUsage.from_bytes(
+        ram_total=1000,
+        ram_available=memory,
+        swap_total=1000,
+        swap_available=1000,
+    )
 
 
-# TODO: this is a hack to get the port for the send_back_multiaddr
-@pytest.fixture
-def create_connection() -> Callable[[NodeId, NodeId, int | None], Connection]:
-    port_counter = 1235
-    ip_counter = 1
+def create_node_network() -> NodeNetworkInfo:
+    return NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address=f"169.254.0.{i}")
+            for i in range(10)
+        ]
+    )
 
-    def _create_connection(
-        source_node_id: NodeId, sink_node_id: NodeId, send_back_port: int | None = None
-    ) -> Connection:
-        nonlocal port_counter
-        nonlocal ip_counter
-        # assign unique ips
-        ip_counter += 1
-        if send_back_port is None:
-            send_back_port = port_counter
-            port_counter += 1
-        return Connection(
-            local_node_id=source_node_id,
-            send_back_node_id=sink_node_id,
-            send_back_multiaddr=Multiaddr(
-                address=f"/ip4/169.254.0.{ip_counter}/tcp/{send_back_port}"
-            ),
-            connection_profile=ConnectionProfile(
-                throughput=1000, latency=1000, jitter=1000
-            ),
-        )
 
-    return _create_connection
+def create_socket_connection(ip: int, sink_port: int = 1234) -> SocketConnection:
+    return SocketConnection(
+        sink_multiaddr=Multiaddr(address=f"/ip4/169.254.0.{ip}/tcp/{sink_port}"),
+    )
+
+
+def create_rdma_connection(iface: int) -> RDMAConnection:
+    return RDMAConnection(
+        source_rdma_iface=f"rdma_en{iface}", sink_rdma_iface=f"rdma_en{iface}"
+    )

src/exo/master/tests/test_api_error_handling.py

@@ -0,0 +1,107 @@
+# pyright: reportUnusedFunction=false, reportAny=false
+from typing import Any, get_args
+
+from fastapi import FastAPI, HTTPException
+from fastapi.testclient import TestClient
+
+from exo.shared.types.api import ErrorInfo, ErrorResponse, FinishReason
+from exo.shared.types.chunks import TokenChunk
+from exo.worker.tests.constants import MODEL_A_ID
+
+
+def test_http_exception_handler_formats_openai_style() -> None:
+    """Test that HTTPException is converted to OpenAI-style error format."""
+    from exo.master.api import API
+
+    app = FastAPI()
+
+    # Setup exception handler
+    api = object.__new__(API)
+    api.app = app
+    api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
+
+    # Add test routes that raise HTTPException
+    @app.get("/test-error")
+    async def _test_error() -> None:
+        raise HTTPException(status_code=500, detail="Test error message")
+
+    @app.get("/test-not-found")
+    async def _test_not_found() -> None:
+        raise HTTPException(status_code=404, detail="Resource not found")
+
+    client = TestClient(app)
+
+    # Test 500 error
+    response = client.get("/test-error")
+    assert response.status_code == 500
+    data: dict[str, Any] = response.json()
+    assert "error" in data
+    assert data["error"]["message"] == "Test error message"
+    assert data["error"]["type"] == "Internal Server Error"
+    assert data["error"]["code"] == 500
+
+    # Test 404 error
+    response = client.get("/test-not-found")
+    assert response.status_code == 404
+    data = response.json()
+    assert "error" in data
+    assert data["error"]["message"] == "Resource not found"
+    assert data["error"]["type"] == "Not Found"
+    assert data["error"]["code"] == 404
+
+
+def test_finish_reason_includes_error() -> None:
+    valid_reasons = get_args(FinishReason)
+    assert "error" in valid_reasons
+
+
+def test_token_chunk_with_error_fields() -> None:
+    chunk = TokenChunk(
+        idx=0,
+        model=MODEL_A_ID,
+        text="",
+        token_id=0,
+        finish_reason="error",
+        error_message="Something went wrong",
+    )
+
+    assert chunk.finish_reason == "error"
+    assert chunk.error_message == "Something went wrong"
+
+
+def test_token_chunk_without_error() -> None:
+    chunk = TokenChunk(
+        idx=1,
+        model=MODEL_A_ID,
+        text="Hello",
+        token_id=42,
+        finish_reason=None,
+    )
+
+    assert chunk.finish_reason is None
+    assert chunk.error_message is None
+
+
+def test_error_response_construction() -> None:
+    error_response = ErrorResponse(
+        error=ErrorInfo(
+            message="Generation failed",
+            type="InternalServerError",
+            code=500,
+        )
+    )
+
+    assert error_response.error.message == "Generation failed"
+    assert error_response.error.code == 500
+
+
+def test_normal_finish_reasons_still_work() -> None:
+    for reason in ["stop", "length", "tool_calls", "content_filter", "function_call"]:
+        chunk = TokenChunk(
+            idx=0,
+            model=MODEL_A_ID,
+            text="done",
+            token_id=100,
+            finish_reason=reason,  # type: ignore[arg-type]
+        )
+        assert chunk.finish_reason == reason

src/exo/master/tests/test_master.py

@@ -7,6 +7,7 @@ from loguru import logger
 
 from exo.master.main import Master
 from exo.routing.router import get_node_id_keypair
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
 from exo.shared.types.commands import (
     ChatCompletion,
@@ -19,15 +20,12 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InstanceCreated,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
     TaskCreated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
+    MemoryUsage,
 )
 from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
 from exo.shared.types.tasks import TaskStatus
@@ -75,29 +73,22 @@ async def test_master():
         tg.start_soon(master.run)
 
         sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
-        # inject a NodePerformanceProfile event
-        logger.info("inject a NodePerformanceProfile event")
+        # inject a NodeGatheredInfo event
+        logger.info("inject a NodeGatheredInfo event")
         await local_event_sender.send(
             ForwarderEvent(
                 origin_idx=0,
                 origin=sender_node_id,
                 session=session_id,
                 event=(
-                    NodePerformanceMeasured(
+                    NodeGatheredInfo(
                         when=str(datetime.now(tz=timezone.utc)),
                         node_id=node_id,
-                        node_profile=NodePerformanceProfile(
-                            model_id="maccy",
-                            chip_id="arm",
-                            friendly_name="test",
-                            memory=MemoryPerformanceProfile(
-                                ram_total=Memory.from_bytes(678948 * 1024),
-                                ram_available=Memory.from_bytes(678948 * 1024),
-                                swap_total=Memory.from_bytes(0),
-                                swap_available=Memory.from_bytes(0),
-                            ),
-                            network_interfaces=[],
-                            system=SystemPerformanceProfile(),
+                        info=MemoryUsage(
+                            ram_total=Memory.from_bytes(678948 * 1024),
+                            ram_available=Memory.from_bytes(678948 * 1024),
+                            swap_total=Memory.from_bytes(0),
+                            swap_available=Memory.from_bytes(0),
                         ),
                     )
                 ),
@@ -108,7 +99,7 @@ async def test_master():
         logger.info("wait for initial topology event")
         while len(list(master.state.topology.list_nodes())) == 0:
             await anyio.sleep(0.001)
-        while len(master.state.node_profiles) == 0:
+        while len(master.state.node_memory) == 0:
             await anyio.sleep(0.001)
 
         logger.info("inject a CreateInstance Command")
@@ -118,9 +109,8 @@ async def test_master():
                 command=(
                     PlaceInstance(
                         command_id=CommandId(),
-                        model_meta=ModelMetadata(
+                        model_card=ModelCard(
                             model_id=ModelId("llama-3.2-1b"),
-                            pretty_name="Llama 3.2 1B",
                             n_layers=16,
                             storage_size=Memory.from_bytes(678948),
                             hidden_size=7168,
@@ -163,7 +153,7 @@ async def test_master():
         assert events[0].idx == 0
         assert events[1].idx == 1
         assert events[2].idx == 2
-        assert isinstance(events[0].event, NodePerformanceMeasured)
+        assert isinstance(events[0].event, NodeGatheredInfo)
         assert isinstance(events[1].event, InstanceCreated)
         created_instance = events[1].event.instance
         assert isinstance(created_instance, MlxRingInstance)
@@ -176,9 +166,8 @@ async def test_master():
                     start_layer=0,
                     end_layer=16,
                     n_layers=16,
-                    model_meta=ModelMetadata(
+                    model_card=ModelCard(
                         model_id=ModelId("llama-3.2-1b"),
-                        pretty_name="Llama 3.2 1B",
                         n_layers=16,
                         storage_size=Memory.from_bytes(678948),
                         hidden_size=7168,

src/exo/master/tests/test_placement.py

@@ -1,20 +1,24 @@
-from typing import Callable
-
 import pytest
-from loguru import logger
 
 from exo.master.placement import (
     get_transition_events,
     place_instance,
 )
+from exo.master.tests.conftest import (
+    create_node_memory,
+    create_node_network,
+    create_rdma_connection,
+    create_socket_connection,
+)
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.topology import Topology
 from exo.shared.types.commands import PlaceInstance
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.events import InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.profiling import NetworkInterfaceInfo, NodePerformanceProfile
-from exo.shared.types.topology import Connection, NodeInfo
+from exo.shared.types.multiaddr import Multiaddr
+from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
+from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.instances import (
     Instance,
     InstanceId,
@@ -26,11 +30,6 @@ from exo.shared.types.worker.runners import ShardAssignments
 from exo.shared.types.worker.shards import Sharding
 
 
-@pytest.fixture
-def topology() -> Topology:
-    return Topology()
-
-
 @pytest.fixture
 def instance() -> Instance:
     return MlxRingInstance(
@@ -44,21 +43,20 @@ def instance() -> Instance:
 
 
 @pytest.fixture
-def model_meta() -> ModelMetadata:
-    return ModelMetadata(
+def model_card() -> ModelCard:
+    return ModelCard(
         model_id=ModelId("test-model"),
         storage_size=Memory.from_kb(1000),
-        pretty_name="Test Model",
         n_layers=10,
         hidden_size=30,
         supports_tensor=True,
     )
 
 
-def place_instance_command(model_meta: ModelMetadata) -> PlaceInstance:
+def place_instance_command(model_card: ModelCard) -> PlaceInstance:
     return PlaceInstance(
         command_id=CommandId(),
-        model_meta=model_meta,
+        model_card=model_card,
         sharding=Sharding.Pipeline,
         instance_meta=InstanceMeta.MlxRing,
         min_nodes=1,
@@ -70,47 +68,75 @@ def place_instance_command(model_meta: ModelMetadata) -> PlaceInstance:
     [
         ((500, 500, 1000), 12, (3, 3, 6)),
         ((500, 500, 500), 12, (4, 4, 4)),
-        ((312, 518, 1024), 12, (2, 3, 7)),
+        ((312, 468, 1092), 12, (2, 3, 7)),
     ],
 )
 def test_get_instance_placements_create_instance(
     available_memory: tuple[int, int, int],
     total_layers: int,
     expected_layers: tuple[int, int, int],
-    topology: Topology,
-    model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
+    model_card: ModelCard,
 ):
     # arrange
-    model_meta.n_layers = total_layers
-    model_meta.storage_size.in_bytes = sum(
+    model_card.n_layers = total_layers
+    model_card.storage_size.in_bytes = sum(
         available_memory
     )  # make it exactly fit across all nodes
+    topology = Topology()
 
-    cic = place_instance_command(model_meta)
+    cic = place_instance_command(model_card)
     node_id_a = NodeId()
     node_id_b = NodeId()
     node_id_c = NodeId()
-    topology.add_node(create_node(available_memory[0], node_id_a))
-    topology.add_node(create_node(available_memory[1], node_id_b))
-    topology.add_node(create_node(available_memory[2], node_id_c))
-    # Add bidirectional connections for ring topology
-    topology.add_connection(create_connection(node_id_a, node_id_b))
-    topology.add_connection(create_connection(node_id_b, node_id_a))
-    topology.add_connection(create_connection(node_id_b, node_id_c))
-    topology.add_connection(create_connection(node_id_c, node_id_b))
-    topology.add_connection(create_connection(node_id_c, node_id_a))
-    topology.add_connection(create_connection(node_id_a, node_id_c))
+
+    # fully connected (directed) between the 3 nodes
+    conn_a_b = Connection(
+        source=node_id_a, sink=node_id_b, edge=create_socket_connection(1)
+    )
+    conn_b_c = Connection(
+        source=node_id_b, sink=node_id_c, edge=create_socket_connection(2)
+    )
+    conn_c_a = Connection(
+        source=node_id_c, sink=node_id_a, edge=create_socket_connection(3)
+    )
+    conn_c_b = Connection(
+        source=node_id_c, sink=node_id_b, edge=create_socket_connection(4)
+    )
+    conn_a_c = Connection(
+        source=node_id_a, sink=node_id_c, edge=create_socket_connection(5)
+    )
+    conn_b_a = Connection(
+        source=node_id_b, sink=node_id_a, edge=create_socket_connection(6)
+    )
+
+    node_memory = {
+        node_id_a: create_node_memory(available_memory[0]),
+        node_id_b: create_node_memory(available_memory[1]),
+        node_id_c: create_node_memory(available_memory[2]),
+    }
+    node_network = {
+        node_id_a: create_node_network(),
+        node_id_b: create_node_network(),
+        node_id_c: create_node_network(),
+    }
+    topology.add_node(node_id_a)
+    topology.add_node(node_id_b)
+    topology.add_node(node_id_c)
+    topology.add_connection(conn_a_b)
+    topology.add_connection(conn_b_c)
+    topology.add_connection(conn_c_a)
+    topology.add_connection(conn_c_b)
+    topology.add_connection(conn_a_c)
+    topology.add_connection(conn_b_a)
 
     # act
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
 
     # assert
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
     instance = placements[instance_id]
-    assert instance.shard_assignments.model_id == model_meta.model_id
+    assert instance.shard_assignments.model_id == model_card.model_id
 
     runner_id_a = instance.shard_assignments.node_to_runner[node_id_a]
     runner_id_b = instance.shard_assignments.node_to_runner[node_id_b]
@@ -130,23 +156,22 @@ def test_get_instance_placements_create_instance(
     assert shards_sorted[-1].end_layer == total_layers
 
 
-def test_get_instance_placements_one_node_exact_fit(
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-) -> None:
+def test_get_instance_placements_one_node_exact_fit() -> None:
     topology = Topology()
     node_id = NodeId()
-    topology.add_node(create_node(1000 * 1024, node_id))
+    topology.add_node(node_id)
+    node_memory = {node_id: create_node_memory(1000 * 1024)}
+    node_network = {node_id: create_node_network()}
     cic = place_instance_command(
-        ModelMetadata(
+        ModelCard(
             model_id=ModelId("test-model"),
             storage_size=Memory.from_kb(1000),
-            pretty_name="Test Model",
             n_layers=10,
             hidden_size=1000,
             supports_tensor=True,
         ),
     )
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -157,23 +182,22 @@ def test_get_instance_placements_one_node_exact_fit(
     assert len(instance.shard_assignments.runner_to_shard) == 1
 
 
-def test_get_instance_placements_one_node_fits_with_extra_memory(
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-) -> None:
+def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
     topology = Topology()
     node_id = NodeId()
-    topology.add_node(create_node(1001 * 1024, node_id))
+    topology.add_node(node_id)
+    node_memory = {node_id: create_node_memory(1001 * 1024)}
+    node_network = {node_id: create_node_network()}
     cic = place_instance_command(
-        ModelMetadata(
+        ModelCard(
             model_id=ModelId("test-model"),
             storage_size=Memory.from_kb(1000),
-            pretty_name="Test Model",
             n_layers=10,
             hidden_size=1000,
             supports_tensor=True,
         ),
     )
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -184,17 +208,16 @@ def test_get_instance_placements_one_node_fits_with_extra_memory(
     assert len(instance.shard_assignments.runner_to_shard) == 1
 
 
-def test_get_instance_placements_one_node_not_fit(
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-) -> None:
+def test_get_instance_placements_one_node_not_fit() -> None:
     topology = Topology()
     node_id = NodeId()
-    topology.add_node(create_node(1000 * 1024, node_id))
+    topology.add_node(node_id)
+    node_memory = {node_id: create_node_memory(1000 * 1024)}
+    node_network = {node_id: create_node_network()}
     cic = place_instance_command(
-        model_meta=ModelMetadata(
+        model_card=ModelCard(
             model_id=ModelId("test-model"),
             storage_size=Memory.from_kb(1001),
-            pretty_name="Test Model",
             n_layers=10,
             hidden_size=1000,
             supports_tensor=True,
@@ -202,7 +225,7 @@ def test_get_instance_placements_one_node_not_fit(
     )
 
     with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
-        place_instance(cic, topology, {})
+        place_instance(cic, topology, {}, node_memory, node_network)
 
 
 def test_get_transition_events_no_change(instance: Instance):
@@ -247,217 +270,175 @@ def test_get_transition_events_delete_instance(instance: Instance):
     assert events[0].instance_id == instance_id
 
 
-def test_placement_selects_cycle_with_most_memory(
-    topology: Topology,
-    model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
+def test_placement_selects_leaf_nodes(
+    model_card: ModelCard,
 ):
-    # Arrange two 3-node cycles with different total memory.
-    # With bidirectional connections for ring topology, both cycles have non-leaf nodes.
-    # The algorithm should select the cycle with the most available memory.
+    # arrange
+    topology = Topology()
 
-    # Model requires more than any single node but fits within a 3-node cycle
-    model_meta.storage_size.in_bytes = 1500
-    model_meta.n_layers = 12
+    model_card.storage_size = Memory.from_bytes(1000)
 
-    # Create node ids
     node_id_a = NodeId()
     node_id_b = NodeId()
     node_id_c = NodeId()
     node_id_d = NodeId()
-    node_id_e = NodeId()
-    node_id_f = NodeId()
 
-    # A-B-C cycle total memory = 1600 (< D-E-F total)
-    topology.add_node(create_node(400, node_id_a))
-    topology.add_node(create_node(400, node_id_b))
-    topology.add_node(create_node(800, node_id_c))
+    node_memory = {
+        node_id_a: create_node_memory(500),
+        node_id_b: create_node_memory(600),
+        node_id_c: create_node_memory(600),
+        node_id_d: create_node_memory(500),
+    }
+    node_network = {
+        node_id_a: create_node_network(),
+        node_id_b: create_node_network(),
+        node_id_c: create_node_network(),
+        node_id_d: create_node_network(),
+    }
 
-    # D-E-F cycle total memory = 1800 (> A-B-C total)
-    topology.add_node(create_node(600, node_id_d))
-    topology.add_node(create_node(600, node_id_e))
-    topology.add_node(create_node(600, node_id_f))
+    topology.add_node(node_id_a)
+    topology.add_node(node_id_b)
+    topology.add_node(node_id_c)
+    topology.add_node(node_id_d)
 
-    # Build bidirectional cycles for ring topology
-    topology.add_connection(create_connection(node_id_a, node_id_b))
-    topology.add_connection(create_connection(node_id_b, node_id_a))
-    topology.add_connection(create_connection(node_id_b, node_id_c))
-    topology.add_connection(create_connection(node_id_c, node_id_b))
-    topology.add_connection(create_connection(node_id_c, node_id_a))
-    topology.add_connection(create_connection(node_id_a, node_id_c))
-
-    topology.add_connection(create_connection(node_id_d, node_id_e))
-    topology.add_connection(create_connection(node_id_e, node_id_d))
-    topology.add_connection(create_connection(node_id_e, node_id_f))
-    topology.add_connection(create_connection(node_id_f, node_id_e))
-    topology.add_connection(create_connection(node_id_f, node_id_d))
-    topology.add_connection(create_connection(node_id_d, node_id_f))
-
-    cic = place_instance_command(
-        model_meta=model_meta,
+    # Daisy chain topology (directed)
+    topology.add_connection(
+        Connection(source=node_id_a, sink=node_id_b, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_b, sink=node_id_a, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_b, sink=node_id_c, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_c, sink=node_id_b, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_c, sink=node_id_d, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_d, sink=node_id_c, edge=create_socket_connection(1))
     )
 
-    # Act
-    placements = place_instance(cic, topology, {})
+    cic = place_instance_command(model_card=model_card)
 
-    # Assert: D-E-F cycle should be selected as it has more total memory
+    # act
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
+
+    # assert
     assert len(placements) == 1
-    instance_id = list(placements.keys())[0]
-    instance = placements[instance_id]
+    instance = list(placements.values())[0]
 
     assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    less_memory_cycle_nodes = {node_id_a, node_id_b, node_id_c}
-    more_memory_cycle_nodes = {node_id_d, node_id_e, node_id_f}
-
-    assert more_memory_cycle_nodes.issubset(assigned_nodes)
-    assert assigned_nodes.isdisjoint(less_memory_cycle_nodes)
+    assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set(
+        (
+            node_id_c,
+            node_id_d,
+        )
+    )
 
 
 def test_tensor_rdma_backend_connectivity_matrix(
-    topology: Topology,
-    model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
+    model_card: ModelCard,
 ):
-    model_meta.n_layers = 12
-    model_meta.storage_size.in_bytes = 1500
+    # arrange
+    topology = Topology()
+    model_card.n_layers = 12
+    model_card.storage_size.in_bytes = 1500
 
-    node_id_a = NodeId()
-    node_id_b = NodeId()
-    node_id_c = NodeId()
+    node_a = NodeId()
+    node_b = NodeId()
+    node_c = NodeId()
 
-    node_a = create_node(500, node_id_a)
-    node_b = create_node(500, node_id_b)
-    node_c = create_node(500, node_id_c)
+    node_memory = {
+        node_a: create_node_memory(500),
+        node_b: create_node_memory(500),
+        node_c: create_node_memory(500),
+    }
 
     ethernet_interface = NetworkInterfaceInfo(
         name="en0",
-        ip_address="192.168.1.100",
+        ip_address="10.0.0.1",
+    )
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
     )
 
-    assert node_a.node_profile is not None
-    assert node_b.node_profile is not None
-    assert node_c.node_profile is not None
-
-    conn_a_b = create_connection(node_id_a, node_id_b)
-    conn_b_c = create_connection(node_id_b, node_id_c)
-    conn_c_a = create_connection(node_id_c, node_id_a)
-
-    conn_b_a = create_connection(node_id_b, node_id_a)
-    conn_c_b = create_connection(node_id_c, node_id_b)
-    conn_a_c = create_connection(node_id_a, node_id_c)
-
-    assert conn_a_b.send_back_multiaddr is not None
-    assert conn_b_c.send_back_multiaddr is not None
-    assert conn_c_a.send_back_multiaddr is not None
-
-    assert conn_b_a.send_back_multiaddr is not None
-    assert conn_c_b.send_back_multiaddr is not None
-    assert conn_a_c.send_back_multiaddr is not None
-
-    node_a.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_a.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_c_a.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_b_a.send_back_multiaddr.ip_address,
-            ),
-            ethernet_interface,
-        ],
-        system=node_a.node_profile.system,
-    )
-    node_b.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_b.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_c_b.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_a_b.send_back_multiaddr.ip_address,
-            ),
-            ethernet_interface,
-        ],
-        system=node_b.node_profile.system,
-    )
-    node_c.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_c.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_a_c.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_b_c.send_back_multiaddr.ip_address,
-            ),
-            ethernet_interface,
-        ],
-        system=node_c.node_profile.system,
-    )
+    node_network = {
+        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
+    }
 
     topology.add_node(node_a)
     topology.add_node(node_b)
     topology.add_node(node_c)
-    topology.add_connection(conn_a_b)
-    topology.add_connection(conn_b_c)
-    topology.add_connection(conn_c_a)
-    topology.add_connection(conn_b_a)
-    topology.add_connection(conn_c_b)
-    topology.add_connection(conn_a_c)
+
+    # RDMA connections (directed)
+    topology.add_connection(
+        Connection(source=node_a, sink=node_b, edge=create_rdma_connection(3))
+    )
+    topology.add_connection(
+        Connection(source=node_b, sink=node_a, edge=create_rdma_connection(3))
+    )
+    topology.add_connection(
+        Connection(source=node_b, sink=node_c, edge=create_rdma_connection(4))
+    )
+    topology.add_connection(
+        Connection(source=node_c, sink=node_b, edge=create_rdma_connection(4))
+    )
+    topology.add_connection(
+        Connection(source=node_a, sink=node_c, edge=create_rdma_connection(5))
+    )
+    topology.add_connection(
+        Connection(source=node_c, sink=node_a, edge=create_rdma_connection(5))
+    )
+
+    # Ethernet connections (directed)
+    topology.add_connection(Connection(source=node_a, sink=node_b, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_b, sink=node_c, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_c, sink=node_a, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_a, sink=node_c, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_b, sink=node_a, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_c, sink=node_b, edge=ethernet_conn))
 
     cic = PlaceInstance(
         sharding=Sharding.Tensor,
         instance_meta=InstanceMeta.MlxJaccl,
         command_id=CommandId(),
-        model_meta=model_meta,
+        model_card=model_card,
         min_nodes=1,
     )
 
-    placements = place_instance(cic, topology, {})
+    # act
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
 
+    # assert
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
     instance = placements[instance_id]
 
     assert isinstance(instance, MlxJacclInstance)
 
-    assert instance.ibv_devices is not None
+    assert instance.jaccl_devices is not None
     assert instance.jaccl_coordinators is not None
 
-    matrix = instance.ibv_devices
+    matrix = instance.jaccl_devices
     assert len(matrix) == 3
-
     for i in range(3):
         assert matrix[i][i] is None
 
     assigned_nodes = list(instance.shard_assignments.node_to_runner.keys())
     node_to_idx = {node_id: idx for idx, node_id in enumerate(assigned_nodes)}
 
-    idx_a = node_to_idx[node_id_a]
-    idx_b = node_to_idx[node_id_b]
-    idx_c = node_to_idx[node_id_c]
+    idx_a = node_to_idx[node_a]
+    idx_b = node_to_idx[node_b]
+    idx_c = node_to_idx[node_c]
 
-    logger.info(matrix)
-
-    assert matrix[idx_a][idx_b] == "rdma_en4"
-    assert matrix[idx_b][idx_c] == "rdma_en3"
-    assert matrix[idx_c][idx_a] == "rdma_en3"
+    assert matrix[idx_a][idx_b] == "rdma_en3"
+    assert matrix[idx_b][idx_c] == "rdma_en4"
+    assert matrix[idx_c][idx_a] == "rdma_en5"
 
     # Verify coordinators are set for all nodes
     assert len(instance.jaccl_coordinators) == 3
@@ -469,7 +450,5 @@ def test_tensor_rdma_backend_connectivity_matrix(
         if node_id == assigned_nodes[0]:
             assert coordinator.startswith("0.0.0.0:")
         else:
-            # Non-rank-0 nodes should have valid IP addresses (can be link-local)
             ip_part = coordinator.split(":")[0]
-            # Just verify it's a valid IP format
             assert len(ip_part.split(".")) == 4

src/exo/master/tests/test_placement_utils.py

@@ -1,162 +1,182 @@
-from typing import Callable
-
 import pytest
 
 from exo.master.placement_utils import (
+    allocate_layers_proportionally,
     filter_cycles_by_memory,
     get_hosts_from_subgraph,
     get_mlx_jaccl_coordinators,
     get_shard_assignments,
     get_smallest_cycles,
 )
+from exo.master.tests.conftest import (
+    create_node_memory,
+    create_socket_connection,
+)
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.profiling import NetworkInterfaceInfo, NodePerformanceProfile
-from exo.shared.types.topology import Connection, NodeInfo
+from exo.shared.types.profiling import (
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
+)
+from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.shards import Sharding
 
 
-@pytest.fixture
-def topology() -> Topology:
-    topology = Topology()
-    return topology
-
-
-def test_filter_cycles_by_memory(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_filter_cycles_by_memory():
     # arrange
     node1_id = NodeId()
     node2_id = NodeId()
+    connection1 = Connection(
+        source=node1_id, sink=node2_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node2_id, sink=node1_id, edge=create_socket_connection(2)
+    )
 
-    node1 = create_node(1000 * 1024, node1_id)
-    node2 = create_node(1000 * 1024, node2_id)
-
-    topology.add_node(node1)
-    topology.add_node(node2)
-
-    connection1 = create_connection(node1_id, node2_id)
-    connection2 = create_connection(node2_id, node1_id)
+    node1_mem = create_node_memory(1000 * 1024)
+    node2_mem = create_node_memory(1000 * 1024)
+    node_memory = {node1_id: node1_mem, node2_id: node2_mem}
 
+    topology = Topology()
+    topology.add_node(node1_id)
+    topology.add_node(node2_id)
     topology.add_connection(connection1)
     topology.add_connection(connection2)
 
-    cycles = topology.get_cycles()
+    cycles = [c for c in topology.get_cycles() if len(c) != 1]
     assert len(cycles) == 1
     assert len(cycles[0]) == 2
 
     # act
-    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_bytes(1))
+    filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_bytes(1))
 
     # assert
     assert len(filtered_cycles) == 1
     assert len(filtered_cycles[0]) == 2
-    assert set(n.node_id for n in filtered_cycles[0]) == {node1_id, node2_id}
+    assert set(n for n in filtered_cycles[0]) == {node1_id, node2_id}
 
 
-def test_filter_cycles_by_insufficient_memory(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_filter_cycles_by_insufficient_memory():
     # arrange
     node1_id = NodeId()
     node2_id = NodeId()
+    connection1 = Connection(
+        source=node1_id, sink=node2_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node2_id, sink=node1_id, edge=create_socket_connection(2)
+    )
 
-    node1 = create_node(1000 * 1024, node1_id)
-    node2 = create_node(1000 * 1024, node2_id)
-
-    topology.add_node(node1)
-    topology.add_node(node2)
-
-    connection1 = create_connection(node1_id, node2_id)
-    connection2 = create_connection(node2_id, node1_id)
+    node1_mem = create_node_memory(1000 * 1024)
+    node2_mem = create_node_memory(1000 * 1024)
+    node_memory = {node1_id: node1_mem, node2_id: node2_mem}
 
+    topology = Topology()
+    topology.add_node(node1_id)
+    topology.add_node(node2_id)
     topology.add_connection(connection1)
     topology.add_connection(connection2)
 
     # act
     filtered_cycles = filter_cycles_by_memory(
-        topology.get_cycles(), Memory.from_kb(2001)
+        topology.get_cycles(), node_memory, Memory.from_kb(2001)
     )
 
     # assert
     assert len(filtered_cycles) == 0
 
 
-def test_filter_multiple_cycles_by_memory(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_filter_multiple_cycles_by_memory():
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_a_id, sink=node_c_id, edge=create_socket_connection(3)
+    )
+    connection4 = Connection(
+        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
+    )
 
-    node_a = create_node(500 * 1024, node_a_id)
-    node_b = create_node(500 * 1024, node_b_id)
-    node_c = create_node(1000 * 1024, node_c_id)
+    node_a_mem = create_node_memory(500 * 1024)
+    node_b_mem = create_node_memory(500 * 1024)
+    node_c_mem = create_node_memory(1000 * 1024)
+    node_memory = {
+        node_a_id: node_a_mem,
+        node_b_id: node_b_mem,
+        node_c_id: node_c_mem,
+    }
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
-
-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
-
-    topology.add_connection(create_connection(node_a_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_b_id))
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
+    topology.add_connection(connection4)
 
     cycles = topology.get_cycles()
 
     # act
-    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_kb(1500))
+    filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_kb(1500))
 
     # assert
     assert len(filtered_cycles) == 1
     assert len(filtered_cycles[0]) == 3
-    assert set(n.node_id for n in filtered_cycles[0]) == {
+    assert set(n for n in filtered_cycles[0]) == {
         node_a_id,
         node_b_id,
         node_c_id,
     }
 
 
-def test_get_smallest_cycles(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_get_smallest_cycles():
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
 
-    node_a = create_node(500 * 1024, node_a_id)
-    node_b = create_node(500 * 1024, node_b_id)
-    node_c = create_node(1000 * 1024, node_c_id)
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_a_id, sink=node_c_id, edge=create_socket_connection(3)
+    )
+    connection4 = Connection(
+        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
+    )
 
-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
+    topology.add_connection(connection4)
+
+    cycles = [c for c in topology.get_cycles() if len(c) != 1]  # ignore singletons
 
     # act
-    smallest_cycles = get_smallest_cycles(topology.get_cycles())
+    smallest_cycles = get_smallest_cycles(cycles)
 
     # assert
     assert len(smallest_cycles) == 1
     assert len(smallest_cycles[0]) == 2
-    assert set(n.node_id for n in smallest_cycles[0]) == {node_a_id, node_b_id}
+    assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id}
 
 
 @pytest.mark.parametrize(
@@ -165,12 +185,12 @@ def test_get_smallest_cycles(
         ((500, 500, 1000), 12, (3, 3, 6)),
         ((500, 500, 500), 12, (4, 4, 4)),
         ((312, 518, 1024), 12, (2, 3, 7)),
+        # Edge case: one node has ~90% of memory - should not over-allocate.
+        # Each node must have enough memory for at least 1 layer (50 KB = 1000/20).
+        ((900, 50, 50), 20, (18, 1, 1)),
     ],
 )
 def test_get_shard_assignments(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
     available_memory: tuple[int, int, int],
     total_layers: int,
     expected_layers: tuple[int, int, int],
@@ -180,44 +200,61 @@ def test_get_shard_assignments(
     node_b_id = NodeId()
     node_c_id = NodeId()
 
-    node_a = create_node(available_memory[0] * 1024, node_a_id)
-    node_b = create_node(available_memory[1] * 1024, node_b_id)
-    node_c = create_node(available_memory[2] * 1024, node_c_id)
+    # create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists)
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
+    connection4 = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(4)
+    )
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
+    topology.add_connection(connection4)
 
-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
+    node_a_mem = create_node_memory(available_memory[0] * 1024)
+    node_b_mem = create_node_memory(available_memory[1] * 1024)
+    node_c_mem = create_node_memory(available_memory[2] * 1024)
+    node_memory = {
+        node_a_id: node_a_mem,
+        node_b_id: node_b_mem,
+        node_c_id: node_c_mem,
+    }
 
-    model_meta = ModelMetadata(
+    model_card = ModelCard(
         model_id=ModelId("test-model"),
-        pretty_name="Test Model",
         n_layers=total_layers,
         storage_size=Memory.from_kb(1000),
         hidden_size=1000,
         supports_tensor=True,
     )
+
     cycles = topology.get_cycles()
-    selected_cycle = cycles[0]
+
+    # pick the 3-node cycle deterministically (cycle ordering can vary)
+    selected_cycle = next(cycle for cycle in cycles if len(cycle) == 3)
 
     # act
     shard_assignments = get_shard_assignments(
-        model_meta, selected_cycle, Sharding.Pipeline
+        model_card, selected_cycle, Sharding.Pipeline, node_memory=node_memory
     )
 
     # assert
     runner_id_a = shard_assignments.node_to_runner[node_a_id]
     runner_id_b = shard_assignments.node_to_runner[node_b_id]
     runner_id_c = shard_assignments.node_to_runner[node_c_id]
-    assert (
-        shard_assignments.runner_to_shard[runner_id_c].end_layer
-        - shard_assignments.runner_to_shard[runner_id_c].start_layer
-        == expected_layers[2]
-    )
+
     assert (
         shard_assignments.runner_to_shard[runner_id_a].end_layer
         - shard_assignments.runner_to_shard[runner_id_a].start_layer
@@ -228,30 +265,37 @@ def test_get_shard_assignments(
         - shard_assignments.runner_to_shard[runner_id_b].start_layer
         == expected_layers[1]
     )
+    assert (
+        shard_assignments.runner_to_shard[runner_id_c].end_layer
+        - shard_assignments.runner_to_shard[runner_id_c].start_layer
+        == expected_layers[2]
+    )
 
 
-def test_get_hosts_from_subgraph(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId, int | None], Connection],
-):
+def test_get_hosts_from_subgraph():
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
+    topology = Topology()
 
-    node_a = create_node(500, node_a_id)
-    node_b = create_node(500, node_b_id)
-    node_c = create_node(1000, node_c_id)
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
 
-    topology.add_connection(create_connection(node_a_id, node_b_id, 5001))
-    topology.add_connection(create_connection(node_b_id, node_c_id, 5002))
-    topology.add_connection(create_connection(node_c_id, node_a_id, 5003))
-    topology.add_connection(create_connection(node_b_id, node_a_id, 5004))
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
 
     # act
     hosts = get_hosts_from_subgraph(topology)
@@ -259,95 +303,68 @@ def test_get_hosts_from_subgraph(
     # assert
     assert len(hosts) == 3
     expected_hosts = [
-        Host(ip=("169.254.0.2"), port=5001),
-        Host(ip=("169.254.0.3"), port=5002),
-        Host(ip=("169.254.0.4"), port=5003),
+        Host(ip="169.254.0.1", port=1234),
+        Host(ip="169.254.0.2", port=1234),
+        Host(ip="169.254.0.3", port=1234),
     ]
     for expected_host in expected_hosts:
         assert expected_host in hosts
 
 
-def test_get_mlx_jaccl_coordinators(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId, int | None], Connection],
-):
+def test_get_mlx_jaccl_coordinators():
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
 
-    node_a = create_node(500 * 1024, node_a_id)
-    node_b = create_node(500 * 1024, node_b_id)
-    node_c = create_node(1000 * 1024, node_c_id)
-
-    conn_a_b = create_connection(node_a_id, node_b_id, 5001)
-    conn_b_a = create_connection(node_b_id, node_a_id, 5002)
-    conn_b_c = create_connection(node_b_id, node_c_id, 5003)
-    conn_c_b = create_connection(node_c_id, node_b_id, 5004)
-    conn_c_a = create_connection(node_c_id, node_a_id, 5005)
-    conn_a_c = create_connection(node_a_id, node_c_id, 5006)
-
-    # Update node profiles with network interfaces before adding to topology
-    assert node_a.node_profile is not None
-    assert node_b.node_profile is not None
-    assert node_c.node_profile is not None
-
-    node_a.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_a.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_a_b.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_a_c.send_back_multiaddr.ip_address,
-            ),
-        ],
-        system=node_a.node_profile.system,
+    # fully connected (directed) between the 3 nodes
+    conn_a_b = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
     )
-    node_b.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_b.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_b_a.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_b_c.send_back_multiaddr.ip_address,
-            ),
-        ],
-        system=node_b.node_profile.system,
+    conn_b_a = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
     )
-    node_c.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_c.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_c_b.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_c_a.send_back_multiaddr.ip_address,
-            ),
-        ],
-        system=node_c.node_profile.system,
+    conn_b_c = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(3)
+    )
+    conn_c_b = Connection(
+        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
+    )
+    conn_c_a = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(5)
+    )
+    conn_a_c = Connection(
+        source=node_a_id, sink=node_c_id, edge=create_socket_connection(6)
     )
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    network_a = NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.5"),
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.2"),
+        ]
+    )
+    network_b = NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.1"),
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.4"),
+        ]
+    )
+    network_c = NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.3"),
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.6"),
+        ]
+    )
+    node_network = {
+        node_a_id: network_a,
+        node_b_id: network_b,
+        node_c_id: network_c,
+    }
+
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
 
     topology.add_connection(conn_a_b)
     topology.add_connection(conn_b_a)
@@ -356,11 +373,12 @@ def test_get_mlx_jaccl_coordinators(
     topology.add_connection(conn_c_a)
     topology.add_connection(conn_a_c)
 
-    cycle = [node_a, node_b, node_c]
-
     # act
     coordinators = get_mlx_jaccl_coordinators(
-        cycle, coordinator_port=5000, cycle_digraph=topology
+        node_a_id,
+        coordinator_port=5000,
+        cycle_digraph=topology,
+        node_network=node_network,
     )
 
     # assert
@@ -381,19 +399,129 @@ def test_get_mlx_jaccl_coordinators(
             f"Coordinator for {node_id} should use port 5000"
         )
 
-    # Rank 0 (node_a) treats this as the listen socket so should listen on all
-    # IPs
+    # Rank 0 (node_a) treats this as the listen socket so should listen on all IPs
     assert coordinators[node_a_id].startswith("0.0.0.0:"), (
-        "Rank 0 node should use localhost as coordinator"
+        "Rank 0 node should use 0.0.0.0 as coordinator listen address"
     )
 
     # Non-rank-0 nodes should use the specific IP from their connection to rank 0
     # node_b uses the IP from conn_b_a (node_b -> node_a)
-    assert coordinators[node_b_id] == (
-        f"{conn_b_a.send_back_multiaddr.ip_address}:5000"
+    assert isinstance(conn_b_a.edge, SocketConnection)
+    assert (
+        coordinators[node_b_id] == f"{conn_b_a.edge.sink_multiaddr.ip_address}:5000"
     ), "node_b should use the IP from conn_b_a"
 
     # node_c uses the IP from conn_c_a (node_c -> node_a)
+    assert isinstance(conn_c_a.edge, SocketConnection)
     assert coordinators[node_c_id] == (
-        f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
+        f"{conn_c_a.edge.sink_multiaddr.ip_address}:5000"
     ), "node_c should use the IP from conn_c_a"
+
+
+class TestAllocateLayersProportionally:
+    def test_empty_node_list_raises(self):
+        with pytest.raises(ValueError, match="empty node list"):
+            allocate_layers_proportionally(total_layers=10, memory_fractions=[])
+
+    def test_zero_layers_raises(self):
+        with pytest.raises(ValueError, match="need at least 1 layer per node"):
+            allocate_layers_proportionally(total_layers=0, memory_fractions=[0.5, 0.5])
+
+    def test_negative_layers_raises(self):
+        with pytest.raises(ValueError, match="need at least 1 layer per node"):
+            allocate_layers_proportionally(total_layers=-1, memory_fractions=[0.5, 0.5])
+
+    def test_fewer_layers_than_nodes_raises(self):
+        with pytest.raises(ValueError, match="need at least 1 layer per node"):
+            allocate_layers_proportionally(
+                total_layers=2, memory_fractions=[0.33, 0.33, 0.34]
+            )
+
+    def test_equal_distribution(self):
+        result = allocate_layers_proportionally(
+            total_layers=12, memory_fractions=[0.25, 0.25, 0.25, 0.25]
+        )
+        assert result == [3, 3, 3, 3]
+        assert sum(result) == 12
+
+    def test_proportional_distribution(self):
+        result = allocate_layers_proportionally(
+            total_layers=12, memory_fractions=[0.25, 0.25, 0.50]
+        )
+        assert result == [3, 3, 6]
+        assert sum(result) == 12
+
+    def test_extreme_imbalance_ensures_minimum(self):
+        result = allocate_layers_proportionally(
+            total_layers=20, memory_fractions=[0.975, 0.0125, 0.0125]
+        )
+        assert all(layers >= 1 for layers in result)
+        assert sum(result) == 20
+        # Small nodes get minimum 1 layer
+        assert result == [18, 1, 1]
+
+    def test_single_node_gets_all_layers(self):
+        result = allocate_layers_proportionally(total_layers=10, memory_fractions=[1.0])
+        assert result == [10]
+
+    def test_minimum_viable_allocation(self):
+        result = allocate_layers_proportionally(
+            total_layers=3, memory_fractions=[0.33, 0.33, 0.34]
+        )
+        assert result == [1, 1, 1]
+        assert sum(result) == 3
+
+
+def test_get_shard_assignments_insufficient_memory_raises():
+    """Test that ValueError is raised when a node has insufficient memory for its layers."""
+    node_a_id = NodeId()
+    node_b_id = NodeId()
+    node_c_id = NodeId()
+    topology = Topology()
+
+    # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers)
+    node_a_mem = create_node_memory(900 * 1024)
+    node_b_mem = create_node_memory(50 * 1024)
+    node_c_mem = create_node_memory(10 * 1024)  # Insufficient memory
+
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
+
+    conn_a_b = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    conn_b_c = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
+    )
+    conn_c_a = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
+    conn_b_a = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
+    topology.add_connection(conn_a_b)
+    topology.add_connection(conn_b_c)
+    topology.add_connection(conn_c_a)
+    topology.add_connection(conn_b_a)
+
+    node_memory = {
+        node_a_id: node_a_mem,
+        node_b_id: node_b_mem,
+        node_c_id: node_c_mem,
+    }
+
+    model_card = ModelCard(
+        model_id=ModelId("test-model"),
+        n_layers=20,
+        storage_size=Memory.from_kb(1000),
+        hidden_size=1000,
+        supports_tensor=True,
+    )
+    cycles = topology.get_cycles()
+    selected_cycle = cycles[0]
+
+    with pytest.raises(ValueError, match="insufficient memory"):
+        get_shard_assignments(
+            model_card, selected_cycle, Sharding.Pipeline, node_memory
+        )

src/exo/master/tests/test_topology.py

@@ -1,13 +1,9 @@
 import pytest
 
 from exo.shared.topology import Topology
+from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
-)
-from exo.shared.types.topology import Connection, ConnectionProfile, NodeId, NodeInfo
+from exo.shared.types.topology import Connection, SocketConnection
 
 
 @pytest.fixture
@@ -16,189 +12,97 @@ def topology() -> Topology:
 
 
 @pytest.fixture
-def connection() -> Connection:
-    return Connection(
-        local_node_id=NodeId(),
-        send_back_node_id=NodeId(),
-        send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
-        connection_profile=ConnectionProfile(
-            throughput=1000, latency=1000, jitter=1000
-        ),
+def socket_connection() -> SocketConnection:
+    return SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
     )
 
 
-@pytest.fixture
-def node_profile() -> NodePerformanceProfile:
-    memory_profile = MemoryPerformanceProfile.from_bytes(
-        ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
-    )
-    system_profile = SystemPerformanceProfile()
-    return NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=memory_profile,
-        network_interfaces=[],
-        system=system_profile,
-    )
-
-
-@pytest.fixture
-def connection_profile() -> ConnectionProfile:
-    return ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
-
-
-def test_add_node(topology: Topology, node_profile: NodePerformanceProfile):
+def test_add_node(topology: Topology):
     # arrange
     node_id = NodeId()
 
     # act
-    topology.add_node(NodeInfo(node_id=node_id, node_profile=node_profile))
+    topology.add_node(node_id)
 
     # assert
-    data = topology.get_node_profile(node_id)
-    assert data == node_profile
+    assert topology.node_is_leaf(node_id)
 
 
-def test_add_connection(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
+def test_add_connection(topology: Topology, socket_connection: SocketConnection):
     # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
+    node_a = NodeId()
+    node_b = NodeId()
+    connection = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
     topology.add_connection(connection)
 
     # act
-    data = topology.get_connection_profile(connection)
+    data = list(topology.list_connections())
 
     # assert
-    assert data == connection.connection_profile
+    assert data == [connection]
 
-
-def test_update_node_profile(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
-    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
-
-    new_node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=MemoryPerformanceProfile.from_bytes(
-            ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
-        ),
-        network_interfaces=[],
-        system=SystemPerformanceProfile(),
-    )
-
-    # act
-    topology.update_node_profile(
-        connection.local_node_id, node_profile=new_node_profile
-    )
-
-    # assert
-    data = topology.get_node_profile(connection.local_node_id)
-    assert data == new_node_profile
-
-
-def test_update_connection_profile(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
-    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
-
-    new_connection_profile = ConnectionProfile(
-        throughput=2000, latency=2000, jitter=2000
-    )
-    connection = Connection(
-        local_node_id=connection.local_node_id,
-        send_back_node_id=connection.send_back_node_id,
-        send_back_multiaddr=connection.send_back_multiaddr,
-        connection_profile=new_connection_profile,
-    )
-
-    # act
-    topology.update_connection_profile(connection)
-
-    # assert
-    data = topology.get_connection_profile(connection)
-    assert data == new_connection_profile
+    assert topology.node_is_leaf(node_a)
+    assert topology.node_is_leaf(node_b)
 
 
 def test_remove_connection_still_connected(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+    topology: Topology, socket_connection: SocketConnection
 ):
     # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
+    node_a = NodeId()
+    node_b = NodeId()
+    conn = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(conn)
 
     # act
-    topology.remove_connection(connection)
+    topology.remove_connection(conn)
 
     # assert
-    assert topology.get_connection_profile(connection) is None
+    assert list(topology.get_all_connections_between(node_a, node_b)) == []
 
 
 def test_remove_node_still_connected(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+    topology: Topology, socket_connection: SocketConnection
 ):
     # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
+    node_a = NodeId()
+    node_b = NodeId()
+    conn = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(conn)
+    assert list(topology.out_edges(node_a)) == [conn]
 
     # act
-    topology.remove_node(connection.local_node_id)
+    topology.remove_node(node_b)
 
     # assert
-    assert topology.get_node_profile(connection.local_node_id) is None
+    assert list(topology.out_edges(node_a)) == []
 
 
-def test_list_nodes(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
+def test_list_nodes(topology: Topology, socket_connection: SocketConnection):
     # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
+    node_a = NodeId()
+    node_b = NodeId()
+    conn = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(conn)
+    assert list(topology.out_edges(node_a)) == [conn]
 
     # act
     nodes = list(topology.list_nodes())
 
     # assert
     assert len(nodes) == 2
-    assert all(isinstance(node, NodeInfo) for node in nodes)
-    assert {node.node_id for node in nodes} == {
-        connection.local_node_id,
-        connection.send_back_node_id,
-    }
+    assert all(isinstance(node, NodeId) for node in nodes)
+    assert set(node for node in nodes) == set([node_a, node_b])

src/exo/shared/apply.py

@@ -11,10 +11,8 @@ from exo.shared.types.events import (
     IndexedEvent,
     InstanceCreated,
     InstanceDeleted,
-    NodeCreated,
     NodeDownloadProgress,
-    NodeMemoryMeasured,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
     NodeTimedOut,
     RunnerDeleted,
     RunnerStatusUpdated,
@@ -27,13 +25,27 @@ from exo.shared.types.events import (
     TopologyEdgeCreated,
     TopologyEdgeDeleted,
 )
-from exo.shared.types.profiling import NodePerformanceProfile, SystemPerformanceProfile
+from exo.shared.types.profiling import (
+    NodeIdentity,
+    NodeNetworkInfo,
+    NodeThunderboltInfo,
+)
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.topology import NodeInfo
+from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
+from exo.utils.info_gatherer.info_gatherer import (
+    MacmonMetrics,
+    MacThunderboltConnections,
+    MacThunderboltIdentifiers,
+    MemoryUsage,
+    MiscData,
+    NodeConfig,
+    NodeNetworkInterfaces,
+    StaticNodeInformation,
+)
 
 
 def event_apply(event: Event, state: State) -> State:
@@ -47,16 +59,12 @@ def event_apply(event: Event, state: State) -> State:
             return apply_instance_created(event, state)
         case InstanceDeleted():
             return apply_instance_deleted(event, state)
-        case NodeCreated():
-            return apply_topology_node_created(event, state)
         case NodeTimedOut():
             return apply_node_timed_out(event, state)
-        case NodePerformanceMeasured():
-            return apply_node_performance_measured(event, state)
         case NodeDownloadProgress():
             return apply_node_download_progress(event, state)
-        case NodeMemoryMeasured():
-            return apply_node_memory_measured(event, state)
+        case NodeGatheredInfo():
+            return apply_node_gathered_info(event, state)
         case RunnerDeleted():
             return apply_runner_deleted(event, state)
         case RunnerStatusUpdated():
@@ -188,120 +196,133 @@ def apply_runner_deleted(event: RunnerDeleted, state: State) -> State:
 
 
 def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
-    topology = copy.copy(state.topology)
-    state.topology.remove_node(event.node_id)
-    node_profiles = {
-        key: value for key, value in state.node_profiles.items() if key != event.node_id
-    }
+    topology = copy.deepcopy(state.topology)
+    topology.remove_node(event.node_id)
     last_seen = {
         key: value for key, value in state.last_seen.items() if key != event.node_id
     }
+    downloads = {
+        key: value for key, value in state.downloads.items() if key != event.node_id
+    }
+    # Clean up all granular node mappings
+    node_identities = {
+        key: value
+        for key, value in state.node_identities.items()
+        if key != event.node_id
+    }
+    node_memory = {
+        key: value for key, value in state.node_memory.items() if key != event.node_id
+    }
+    node_system = {
+        key: value for key, value in state.node_system.items() if key != event.node_id
+    }
+    node_network = {
+        key: value for key, value in state.node_network.items() if key != event.node_id
+    }
+    node_thunderbolt = {
+        key: value
+        for key, value in state.node_thunderbolt.items()
+        if key != event.node_id
+    }
     return state.model_copy(
         update={
+            "downloads": downloads,
             "topology": topology,
-            "node_profiles": node_profiles,
             "last_seen": last_seen,
+            "node_identities": node_identities,
+            "node_memory": node_memory,
+            "node_system": node_system,
+            "node_network": node_network,
+            "node_thunderbolt": node_thunderbolt,
         }
     )
 
 
-def apply_node_performance_measured(
-    event: NodePerformanceMeasured, state: State
-) -> State:
-    new_profiles: Mapping[NodeId, NodePerformanceProfile] = {
-        **state.node_profiles,
-        event.node_id: event.node_profile,
-    }
-    last_seen: Mapping[NodeId, datetime] = {
-        **state.last_seen,
-        event.node_id: datetime.fromisoformat(event.when),
-    }
-    state = state.model_copy(update={"node_profiles": new_profiles})
-    topology = copy.copy(state.topology)
-    # TODO: NodeCreated
-    if not topology.contains_node(event.node_id):
-        topology.add_node(NodeInfo(node_id=event.node_id))
-    topology.update_node_profile(event.node_id, event.node_profile)
-    return state.model_copy(
-        update={
-            "node_profiles": new_profiles,
-            "topology": topology,
-            "last_seen": last_seen,
-        }
-    )
+def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
+    topology = copy.deepcopy(state.topology)
+    topology.add_node(event.node_id)
+    info = event.info
 
-
-def apply_node_memory_measured(event: NodeMemoryMeasured, state: State) -> State:
-    existing = state.node_profiles.get(event.node_id)
-    topology = copy.copy(state.topology)
-
-    if existing is None:
-        created = NodePerformanceProfile(
-            model_id="unknown",
-            chip_id="unknown",
-            friendly_name="Unknown",
-            memory=event.memory,
-            network_interfaces=[],
-            system=SystemPerformanceProfile(
-                # TODO: flops_fp16=0.0,
-                gpu_usage=0.0,
-                temp=0.0,
-                sys_power=0.0,
-                pcpu_usage=0.0,
-                ecpu_usage=0.0,
-                ane_power=0.0,
-            ),
-        )
-        created_profiles: Mapping[NodeId, NodePerformanceProfile] = {
-            **state.node_profiles,
-            event.node_id: created,
-        }
-        last_seen: Mapping[NodeId, datetime] = {
+    # Build update dict with only the mappings that change
+    update: dict[str, object] = {
+        "last_seen": {
             **state.last_seen,
             event.node_id: datetime.fromisoformat(event.when),
-        }
-        if not topology.contains_node(event.node_id):
-            topology.add_node(NodeInfo(node_id=event.node_id))
-            # TODO: NodeCreated
-        topology.update_node_profile(event.node_id, created)
-        return state.model_copy(
-            update={
-                "node_profiles": created_profiles,
-                "topology": topology,
-                "last_seen": last_seen,
-            }
-        )
-
-    updated = existing.model_copy(update={"memory": event.memory})
-    updated_profiles: Mapping[NodeId, NodePerformanceProfile] = {
-        **state.node_profiles,
-        event.node_id: updated,
+        },
+        "topology": topology,
     }
-    # TODO: NodeCreated
-    if not topology.contains_node(event.node_id):
-        topology.add_node(NodeInfo(node_id=event.node_id))
-    topology.update_node_profile(event.node_id, updated)
-    return state.model_copy(
-        update={"node_profiles": updated_profiles, "topology": topology}
-    )
 
+    match info:
+        case MacmonMetrics():
+            update["node_system"] = {
+                **state.node_system,
+                event.node_id: info.system_profile,
+            }
+            update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
+        case MemoryUsage():
+            update["node_memory"] = {**state.node_memory, event.node_id: info}
+        case NodeConfig():
+            pass
+        case MiscData():
+            current_identity = state.node_identities.get(event.node_id, NodeIdentity())
+            new_identity = current_identity.model_copy(
+                update={"friendly_name": info.friendly_name}
+            )
+            update["node_identities"] = {
+                **state.node_identities,
+                event.node_id: new_identity,
+            }
+        case StaticNodeInformation():
+            current_identity = state.node_identities.get(event.node_id, NodeIdentity())
+            new_identity = current_identity.model_copy(
+                update={"model_id": info.model, "chip_id": info.chip}
+            )
+            update["node_identities"] = {
+                **state.node_identities,
+                event.node_id: new_identity,
+            }
+        case NodeNetworkInterfaces():
+            update["node_network"] = {
+                **state.node_network,
+                event.node_id: NodeNetworkInfo(interfaces=info.ifaces),
+            }
+        case MacThunderboltIdentifiers():
+            update["node_thunderbolt"] = {
+                **state.node_thunderbolt,
+                event.node_id: NodeThunderboltInfo(interfaces=info.idents),
+            }
+        case MacThunderboltConnections():
+            conn_map = {
+                tb_ident.domain_uuid: (nid, tb_ident.rdma_interface)
+                for nid in state.node_thunderbolt
+                for tb_ident in state.node_thunderbolt[nid].interfaces
+            }
+            as_rdma_conns = [
+                Connection(
+                    source=event.node_id,
+                    sink=conn_map[tb_conn.sink_uuid][0],
+                    edge=RDMAConnection(
+                        source_rdma_iface=conn_map[tb_conn.source_uuid][1],
+                        sink_rdma_iface=conn_map[tb_conn.sink_uuid][1],
+                    ),
+                )
+                for tb_conn in info.conns
+                if tb_conn.source_uuid in conn_map
+                if tb_conn.sink_uuid in conn_map
+            ]
+            topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
 
-def apply_topology_node_created(event: NodeCreated, state: State) -> State:
-    topology = copy.copy(state.topology)
-    topology.add_node(NodeInfo(node_id=event.node_id))
-    return state.model_copy(update={"topology": topology})
+    return state.model_copy(update=update)
 
 
 def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
-    topology = copy.copy(state.topology)
-    topology.add_connection(event.edge)
+    topology = copy.deepcopy(state.topology)
+    topology.add_connection(event.conn)
     return state.model_copy(update={"topology": topology})
 
 
 def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
-    topology = copy.copy(state.topology)
-    if not topology.contains_connection(event.edge):
-        return state
-    topology.remove_connection(event.edge)
+    topology = copy.deepcopy(state.topology)
+    topology.remove_connection(event.conn)
     # TODO: Clean up removing the reverse connection
     return state.model_copy(update={"topology": topology})

src/exo/shared/constants.py

@@ -38,6 +38,7 @@ EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
 
 # Identity (config)
 EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"
+EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml"
 
 # libp2p topics for event forwarding
 LIBP2P_LOCAL_EVENTS_TOPIC = "worker_events"

src/exo/shared/logging.py

@@ -11,9 +11,6 @@ class InterceptLogger(HypercornLogger):
     def __init__(self, config: Config):
         super().__init__(config)
         assert self.error_logger
-        # TODO: Decide if we want to provide access logs
-        # assert self.access_logger
-        # self.access_logger.handlers = [_InterceptHandler()]
         self.error_logger.handlers = [_InterceptHandler()]
 
 
@@ -29,6 +26,11 @@ class _InterceptHandler(logging.Handler):
 
 def logger_setup(log_file: Path | None, verbosity: int = 0):
     """Set up logging for this process - formatting, file handles, verbosity and output"""
+
+    logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    logging.getLogger("httpcore").setLevel(logging.WARNING)
+
     logger.remove()
 
     # replace all stdlib loggers with _InterceptHandlers that log to loguru

src/exo/shared/models/model_cards.py

@@ -1,552 +1,310 @@
+from pydantic import PositiveInt
+
+from exo.shared.types.common import Id
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.utils.pydantic_ext import CamelCaseModel
 
 
+class ModelId(Id):
+    def normalize(self) -> str:
+        return self.replace("/", "--")
+
+    def short(self) -> str:
+        return self.split("/")[-1]
+
+
 class ModelCard(CamelCaseModel):
-    short_id: str
     model_id: ModelId
-    name: str
-    description: str
-    tags: list[str]
-    metadata: ModelMetadata
+    storage_size: Memory
+    n_layers: PositiveInt
+    hidden_size: PositiveInt
+    supports_tensor: bool
 
 
 MODEL_CARDS: dict[str, ModelCard] = {
     # deepseek v3
     "deepseek-v3.1-4bit": ModelCard(
-        short_id="deepseek-v3.1-4bit",
         model_id=ModelId("mlx-community/DeepSeek-V3.1-4bit"),
-        name="DeepSeek V3.1 (4-bit)",
-        description="""DeepSeek V3.1 is a large language model trained on the DeepSeek V3.1 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/DeepSeek-V3.1-4bit"),
-            pretty_name="DeepSeek V3.1 (4-bit)",
-            storage_size=Memory.from_gb(378),
-            n_layers=61,
-            hidden_size=7168,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(378),
+        n_layers=61,
+        hidden_size=7168,
+        supports_tensor=True,
     ),
     "deepseek-v3.1-8bit": ModelCard(
-        short_id="deepseek-v3.1-8bit",
         model_id=ModelId("mlx-community/DeepSeek-V3.1-8bit"),
-        name="DeepSeek V3.1 (8-bit)",
-        description="""DeepSeek V3.1 is a large language model trained on the DeepSeek V3.1 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/DeepSeek-V3.1-8bit"),
-            pretty_name="DeepSeek V3.1 (8-bit)",
-            storage_size=Memory.from_gb(713),
-            n_layers=61,
-            hidden_size=7168,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(713),
+        n_layers=61,
+        hidden_size=7168,
+        supports_tensor=True,
     ),
     # kimi k2
     "kimi-k2-instruct-4bit": ModelCard(
-        short_id="kimi-k2-instruct-4bit",
         model_id=ModelId("mlx-community/Kimi-K2-Instruct-4bit"),
-        name="Kimi K2 Instruct (4-bit)",
-        description="""Kimi K2 is a large language model trained on the Kimi K2 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Kimi-K2-Instruct-4bit"),
-            pretty_name="Kimi K2 Instruct (4-bit)",
-            storage_size=Memory.from_gb(578),
-            n_layers=61,
-            hidden_size=7168,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(578),
+        n_layers=61,
+        hidden_size=7168,
+        supports_tensor=True,
     ),
     "kimi-k2-thinking": ModelCard(
-        short_id="kimi-k2-thinking",
         model_id=ModelId("mlx-community/Kimi-K2-Thinking"),
-        name="Kimi K2 Thinking (4-bit)",
-        description="""Kimi K2 Thinking is the latest, most capable version of open-source thinking model.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Kimi-K2-Thinking"),
-            pretty_name="Kimi K2 Thinking (4-bit)",
-            storage_size=Memory.from_gb(658),
-            n_layers=61,
-            hidden_size=7168,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(658),
+        n_layers=61,
+        hidden_size=7168,
+        supports_tensor=True,
     ),
     # llama-3.1
     "llama-3.1-8b": ModelCard(
-        short_id="llama-3.1-8b",
         model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"),
-        name="Llama 3.1 8B (4-bit)",
-        description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"),
-            pretty_name="Llama 3.1 8B (4-bit)",
-            storage_size=Memory.from_mb(4423),
-            n_layers=32,
-            hidden_size=4096,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(4423),
+        n_layers=32,
+        hidden_size=4096,
+        supports_tensor=True,
     ),
     "llama-3.1-8b-8bit": ModelCard(
-        short_id="llama-3.1-8b-8bit",
         model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"),
-        name="Llama 3.1 8B (8-bit)",
-        description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"),
-            pretty_name="Llama 3.1 8B (8-bit)",
-            storage_size=Memory.from_mb(8540),
-            n_layers=32,
-            hidden_size=4096,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(8540),
+        n_layers=32,
+        hidden_size=4096,
+        supports_tensor=True,
     ),
     "llama-3.1-8b-bf16": ModelCard(
-        short_id="llama-3.1-8b-bf16",
         model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"),
-        name="Llama 3.1 8B (BF16)",
-        description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"),
-            pretty_name="Llama 3.1 8B (BF16)",
-            storage_size=Memory.from_mb(16100),
-            n_layers=32,
-            hidden_size=4096,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(16100),
+        n_layers=32,
+        hidden_size=4096,
+        supports_tensor=True,
     ),
     "llama-3.1-70b": ModelCard(
-        short_id="llama-3.1-70b",
         model_id=ModelId("mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"),
-        name="Llama 3.1 70B (4-bit)",
-        description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"),
-            pretty_name="Llama 3.1 70B (4-bit)",
-            storage_size=Memory.from_mb(38769),
-            n_layers=80,
-            hidden_size=8192,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(38769),
+        n_layers=80,
+        hidden_size=8192,
+        supports_tensor=True,
     ),
     # llama-3.2
     "llama-3.2-1b": ModelCard(
-        short_id="llama-3.2-1b",
         model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
-        name="Llama 3.2 1B (4-bit)",
-        description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
-            pretty_name="Llama 3.2 1B (4-bit)",
-            storage_size=Memory.from_mb(696),
-            n_layers=16,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(696),
+        n_layers=16,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "llama-3.2-3b": ModelCard(
-        short_id="llama-3.2-3b",
         model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-4bit"),
-        name="Llama 3.2 3B (4-bit)",
-        description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-4bit"),
-            pretty_name="Llama 3.2 3B (4-bit)",
-            storage_size=Memory.from_mb(1777),
-            n_layers=28,
-            hidden_size=3072,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(1777),
+        n_layers=28,
+        hidden_size=3072,
+        supports_tensor=True,
     ),
     "llama-3.2-3b-8bit": ModelCard(
-        short_id="llama-3.2-3b-8bit",
         model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-8bit"),
-        name="Llama 3.2 3B (8-bit)",
-        description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-8bit"),
-            pretty_name="Llama 3.2 3B (8-bit)",
-            storage_size=Memory.from_mb(3339),
-            n_layers=28,
-            hidden_size=3072,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(3339),
+        n_layers=28,
+        hidden_size=3072,
+        supports_tensor=True,
     ),
     # llama-3.3
     "llama-3.3-70b": ModelCard(
-        short_id="llama-3.3-70b",
         model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-4bit"),
-        name="Llama 3.3 70B (4-bit)",
-        description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-4bit"),
-            pretty_name="Llama 3.3 70B",
-            storage_size=Memory.from_mb(38769),
-            n_layers=80,
-            hidden_size=8192,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(38769),
+        n_layers=80,
+        hidden_size=8192,
+        supports_tensor=True,
     ),
     "llama-3.3-70b-8bit": ModelCard(
-        short_id="llama-3.3-70b-8bit",
         model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-8bit"),
-        name="Llama 3.3 70B (8-bit)",
-        description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-8bit"),
-            pretty_name="Llama 3.3 70B (8-bit)",
-            storage_size=Memory.from_mb(73242),
-            n_layers=80,
-            hidden_size=8192,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(73242),
+        n_layers=80,
+        hidden_size=8192,
+        supports_tensor=True,
     ),
     "llama-3.3-70b-fp16": ModelCard(
-        short_id="llama-3.3-70b-fp16",
         model_id=ModelId("mlx-community/llama-3.3-70b-instruct-fp16"),
-        name="Llama 3.3 70B (FP16)",
-        description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/llama-3.3-70b-instruct-fp16"),
-            pretty_name="Llama 3.3 70B (FP16)",
-            storage_size=Memory.from_mb(137695),
-            n_layers=80,
-            hidden_size=8192,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(137695),
+        n_layers=80,
+        hidden_size=8192,
+        supports_tensor=True,
     ),
     # qwen3
     "qwen3-0.6b": ModelCard(
-        short_id="qwen3-0.6b",
         model_id=ModelId("mlx-community/Qwen3-0.6B-4bit"),
-        name="Qwen3 0.6B (4-bit)",
-        description="""Qwen3 0.6B is a large language model trained on the Qwen3 0.6B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-0.6B-4bit"),
-            pretty_name="Qwen3 0.6B (4-bit)",
-            storage_size=Memory.from_mb(327),
-            n_layers=28,
-            hidden_size=1024,
-            supports_tensor=False,
-        ),
+        storage_size=Memory.from_mb(327),
+        n_layers=28,
+        hidden_size=1024,
+        supports_tensor=False,
     ),
     "qwen3-0.6b-8bit": ModelCard(
-        short_id="qwen3-0.6b-8bit",
         model_id=ModelId("mlx-community/Qwen3-0.6B-8bit"),
-        name="Qwen3 0.6B (8-bit)",
-        description="""Qwen3 0.6B is a large language model trained on the Qwen3 0.6B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-0.6B-8bit"),
-            pretty_name="Qwen3 0.6B (8-bit)",
-            storage_size=Memory.from_mb(666),
-            n_layers=28,
-            hidden_size=1024,
-            supports_tensor=False,
-        ),
+        storage_size=Memory.from_mb(666),
+        n_layers=28,
+        hidden_size=1024,
+        supports_tensor=False,
     ),
     "qwen3-30b": ModelCard(
-        short_id="qwen3-30b",
         model_id=ModelId("mlx-community/Qwen3-30B-A3B-4bit"),
-        name="Qwen3 30B A3B (4-bit)",
-        description="""Qwen3 30B is a large language model trained on the Qwen3 30B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-30B-A3B-4bit"),
-            pretty_name="Qwen3 30B A3B (4-bit)",
-            storage_size=Memory.from_mb(16797),
-            n_layers=48,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(16797),
+        n_layers=48,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "qwen3-30b-8bit": ModelCard(
-        short_id="qwen3-30b-8bit",
         model_id=ModelId("mlx-community/Qwen3-30B-A3B-8bit"),
-        name="Qwen3 30B A3B (8-bit)",
-        description="""Qwen3 30B is a large language model trained on the Qwen3 30B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-30B-A3B-8bit"),
-            pretty_name="Qwen3 30B A3B (8-bit)",
-            storage_size=Memory.from_mb(31738),
-            n_layers=48,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(31738),
+        n_layers=48,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "qwen3-80b-a3B-4bit": ModelCard(
-        short_id="qwen3-80b-a3B-4bit",
         model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"),
-        name="Qwen3 80B A3B (4-bit)",
-        description="""Qwen3 80B""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"),
-            pretty_name="Qwen3 80B A3B (4-bit)",
-            storage_size=Memory.from_mb(44800),
-            n_layers=48,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(44800),
+        n_layers=48,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "qwen3-80b-a3B-8bit": ModelCard(
-        short_id="qwen3-80b-a3B-8bit",
         model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"),
-        name="Qwen3 80B A3B (8-bit)",
-        description="""Qwen3 80B""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"),
-            pretty_name="Qwen3 80B A3B (8-bit)",
-            storage_size=Memory.from_mb(84700),
-            n_layers=48,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(84700),
+        n_layers=48,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "qwen3-80b-a3B-thinking-4bit": ModelCard(
-        short_id="qwen3-80b-a3B-thinking-4bit",
         model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
-        name="Qwen3 80B A3B Thinking (4-bit)",
-        description="""Qwen3 80B Reasoning model""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
-            pretty_name="Qwen3 80B A3B (4-bit)",
-            storage_size=Memory.from_mb(84700),
-            n_layers=48,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(84700),
+        n_layers=48,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "qwen3-80b-a3B-thinking-8bit": ModelCard(
-        short_id="qwen3-80b-a3B-thinking-8bit",
         model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"),
-        name="Qwen3 80B A3B Thinking (8-bit)",
-        description="""Qwen3 80B Reasoning model""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"),
-            pretty_name="Qwen3 80B A3B (8-bit)",
-            storage_size=Memory.from_mb(84700),
-            n_layers=48,
-            hidden_size=2048,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_mb(84700),
+        n_layers=48,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     "qwen3-235b-a22b-4bit": ModelCard(
-        short_id="qwen3-235b-a22b-4bit",
         model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"),
-        name="Qwen3 235B A22B (4-bit)",
-        description="""Qwen3 235B (Active 22B) is a large language model trained on the Qwen3 235B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"),
-            pretty_name="Qwen3 235B A22B (4-bit)",
-            storage_size=Memory.from_gb(132),
-            n_layers=94,
-            hidden_size=4096,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(132),
+        n_layers=94,
+        hidden_size=4096,
+        supports_tensor=True,
     ),
     "qwen3-235b-a22b-8bit": ModelCard(
-        short_id="qwen3-235b-a22b-8bit",
         model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"),
-        name="Qwen3 235B A22B (8-bit)",
-        description="""Qwen3 235B (Active 22B) is a large language model trained on the Qwen3 235B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"),
-            pretty_name="Qwen3 235B A22B (8-bit)",
-            storage_size=Memory.from_gb(250),
-            n_layers=94,
-            hidden_size=4096,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(250),
+        n_layers=94,
+        hidden_size=4096,
+        supports_tensor=True,
     ),
     "qwen3-coder-480b-a35b-4bit": ModelCard(
-        short_id="qwen3-coder-480b-a35b-4bit",
         model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"),
-        name="Qwen3 Coder 480B A35B (4-bit)",
-        description="""Qwen3 Coder 480B (Active 35B) is a large language model trained on the Qwen3 Coder 480B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"),
-            pretty_name="Qwen3 Coder 480B A35B (4-bit)",
-            storage_size=Memory.from_gb(270),
-            n_layers=62,
-            hidden_size=6144,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(270),
+        n_layers=62,
+        hidden_size=6144,
+        supports_tensor=True,
     ),
     "qwen3-coder-480b-a35b-8bit": ModelCard(
-        short_id="qwen3-coder-480b-a35b-8bit",
         model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"),
-        name="Qwen3 Coder 480B A35B (8-bit)",
-        description="""Qwen3 Coder 480B (Active 35B) is a large language model trained on the Qwen3 Coder 480B dataset.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"),
-            pretty_name="Qwen3 Coder 480B A35B (8-bit)",
-            storage_size=Memory.from_gb(540),
-            n_layers=62,
-            hidden_size=6144,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(540),
+        n_layers=62,
+        hidden_size=6144,
+        supports_tensor=True,
     ),
     # gpt-oss
     "gpt-oss-120b-MXFP4-Q8": ModelCard(
-        short_id="gpt-oss-120b-MXFP4-Q8",
         model_id=ModelId("mlx-community/gpt-oss-120b-MXFP4-Q8"),
-        name="GPT-OSS 120B (MXFP4-Q8, MLX)",
-        description="""OpenAI's GPT-OSS 120B is a 117B-parameter Mixture-of-Experts model designed for high-reasoning and general-purpose use; this variant is a 4-bit MLX conversion for Apple Silicon.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/gpt-oss-120b-MXFP4-Q8"),
-            pretty_name="GPT-OSS 120B (MXFP4-Q8, MLX)",
-            storage_size=Memory.from_kb(68_996_301),
-            n_layers=36,
-            hidden_size=2880,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_kb(68_996_301),
+        n_layers=36,
+        hidden_size=2880,
+        supports_tensor=True,
     ),
     "gpt-oss-20b-MXFP4-Q8": ModelCard(
-        short_id="gpt-oss-20b-MXFP4-Q8",
         model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
-        name="GPT-OSS 20B (MXFP4-Q8, MLX)",
-        description="""OpenAI's GPT-OSS 20B is a medium-sized MoE model for lower-latency and local or specialized use cases; this variant is a 4-bit MLX conversion for Apple Silicon.""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
-            pretty_name="GPT-OSS 20B (MXFP4-Q8, MLX)",
-            storage_size=Memory.from_kb(11_744_051),
-            n_layers=24,
-            hidden_size=2880,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_kb(11_744_051),
+        n_layers=24,
+        hidden_size=2880,
+        supports_tensor=True,
     ),
     # glm 4.5
     "glm-4.5-air-8bit": ModelCard(
         # Needs to be quantized g32 or g16 to work with tensor parallel
-        short_id="glm-4.5-air-8bit",
         model_id=ModelId("mlx-community/GLM-4.5-Air-8bit"),
-        name="GLM 4.5 Air 8bit",
-        description="""GLM 4.5 Air 8bit""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/GLM-4.5-Air-8bit"),
-            pretty_name="GLM 4.5 Air 8bit",
-            storage_size=Memory.from_gb(114),
-            n_layers=46,
-            hidden_size=4096,
-            supports_tensor=False,
-        ),
+        storage_size=Memory.from_gb(114),
+        n_layers=46,
+        hidden_size=4096,
+        supports_tensor=False,
     ),
     "glm-4.5-air-bf16": ModelCard(
-        short_id="glm-4.5-air-bf16",
         model_id=ModelId("mlx-community/GLM-4.5-Air-bf16"),
-        name="GLM 4.5 Air bf16",
-        description="""GLM 4.5 Air bf16""",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/GLM-4.5-Air-bf16"),
-            pretty_name="GLM 4.5 Air bf16",
-            storage_size=Memory.from_gb(214),
-            n_layers=46,
-            hidden_size=4096,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_gb(214),
+        n_layers=46,
+        hidden_size=4096,
+        supports_tensor=True,
     ),
     # glm 4.7
     "glm-4.7-4bit": ModelCard(
-        short_id="glm-4.7-4bit",
         model_id=ModelId("mlx-community/GLM-4.7-4bit"),
-        name="GLM 4.7 4bit",
-        description="GLM 4.7 4bit",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/GLM-4.7-4bit"),
-            pretty_name="GLM 4.7 4bit",
-            storage_size=Memory.from_bytes(198556925568),
-            n_layers=91,
-            hidden_size=5120,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_bytes(198556925568),
+        n_layers=91,
+        hidden_size=5120,
+        supports_tensor=True,
     ),
     "glm-4.7-6bit": ModelCard(
-        short_id="glm-4.7-6bit",
         model_id=ModelId("mlx-community/GLM-4.7-6bit"),
-        name="GLM 4.7 6bit",
-        description="GLM 4.7 6bit",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/GLM-4.7-6bit"),
-            pretty_name="GLM 4.7 6bit",
-            storage_size=Memory.from_bytes(286737579648),
-            n_layers=91,
-            hidden_size=5120,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_bytes(286737579648),
+        n_layers=91,
+        hidden_size=5120,
+        supports_tensor=True,
     ),
     "glm-4.7-8bit-gs32": ModelCard(
-        short_id="glm-4.7-8bit-gs32",
         model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"),
-        name="GLM 4.7 8bit (gs32)",
-        description="GLM 4.7 8bit (gs32)",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"),
-            pretty_name="GLM 4.7 8bit (gs32)",
-            storage_size=Memory.from_bytes(396963397248),
-            n_layers=91,
-            hidden_size=5120,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_bytes(396963397248),
+        n_layers=91,
+        hidden_size=5120,
+        supports_tensor=True,
+    ),
+    # glm 4.7 flash
+    "glm-4.7-flash-4bit": ModelCard(
+        model_id=ModelId("mlx-community/GLM-4.7-Flash-4bit"),
+        storage_size=Memory.from_gb(18),
+        n_layers=47,
+        hidden_size=2048,
+        supports_tensor=True,
+    ),
+    "glm-4.7-flash-5bit": ModelCard(
+        model_id=ModelId("mlx-community/GLM-4.7-Flash-5bit"),
+        storage_size=Memory.from_gb(21),
+        n_layers=47,
+        hidden_size=2048,
+        supports_tensor=True,
+    ),
+    "glm-4.7-flash-6bit": ModelCard(
+        model_id=ModelId("mlx-community/GLM-4.7-Flash-6bit"),
+        storage_size=Memory.from_gb(25),
+        n_layers=47,
+        hidden_size=2048,
+        supports_tensor=True,
+    ),
+    "glm-4.7-flash-8bit": ModelCard(
+        model_id=ModelId("mlx-community/GLM-4.7-Flash-8bit"),
+        storage_size=Memory.from_gb(32),
+        n_layers=47,
+        hidden_size=2048,
+        supports_tensor=True,
     ),
     # minimax-m2
     "minimax-m2.1-8bit": ModelCard(
-        short_id="minimax-m2.1-8bit",
         model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"),
-        name="MiniMax M2.1 8bit",
-        description="MiniMax M2.1 8bit",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"),
-            pretty_name="MiniMax M2.1 8bit",
-            storage_size=Memory.from_bytes(242986745856),
-            n_layers=61,
-            hidden_size=3072,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_bytes(242986745856),
+        n_layers=61,
+        hidden_size=3072,
+        supports_tensor=True,
     ),
     "minimax-m2.1-3bit": ModelCard(
-        short_id="minimax-m2.1-3bit",
         model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"),
-        name="MiniMax M2.1 3bit",
-        description="MiniMax M2.1 3bit",
-        tags=[],
-        metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"),
-            pretty_name="MiniMax M2.1 3bit",
-            storage_size=Memory.from_bytes(100086644736),
-            n_layers=61,
-            hidden_size=3072,
-            supports_tensor=True,
-        ),
+        storage_size=Memory.from_bytes(100086644736),
+        n_layers=61,
+        hidden_size=3072,
+        supports_tensor=True,
     ),
 }

src/exo/shared/models/model_meta.py

@@ -6,9 +6,8 @@ from huggingface_hub import model_info
 from loguru import logger
 from pydantic import BaseModel, Field
 
-from exo.shared.models.model_cards import MODEL_CARDS
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.worker.download.download_utils import (
     ModelSafetensorsIndex,
     download_file_with_retry,
@@ -92,18 +91,18 @@ async def get_safetensors_size(model_id: str) -> Memory:
     return Memory.from_bytes(info.safetensors.total)
 
 
-_model_meta_cache: dict[str, ModelMetadata] = {}
+_model_card_cache: dict[str, ModelCard] = {}
 
 
-async def get_model_meta(model_id: str) -> ModelMetadata:
-    if model_id in _model_meta_cache:
-        return _model_meta_cache[model_id]
-    model_meta = await _get_model_meta(model_id)
-    _model_meta_cache[model_id] = model_meta
-    return model_meta
+async def get_model_card(model_id: str) -> ModelCard:
+    if model_id in _model_card_cache:
+        return _model_card_cache[model_id]
+    model_card = await _get_model_card(model_id)
+    _model_card_cache[model_id] = model_card
+    return model_card
 
 
-async def _get_model_meta(model_id: str) -> ModelMetadata:
+async def _get_model_card(model_id: str) -> ModelCard:
     """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
     config_data = await get_config_data(model_id)
     num_layers = config_data.layer_count
@@ -113,14 +112,11 @@ async def _get_model_meta(model_id: str) -> ModelMetadata:
         None,
     )
 
-    return ModelMetadata(
+    return ModelCard(
         model_id=ModelId(model_id),
-        pretty_name=model_card.name if model_card is not None else model_id,
         storage_size=mem_size_bytes,
         n_layers=num_layers,
         hidden_size=config_data.hidden_size or 0,
         # TODO: all custom models currently do not support tensor. We could add a dynamic test for this?
-        supports_tensor=model_card.metadata.supports_tensor
-        if model_card is not None
-        else False,
+        supports_tensor=model_card.supports_tensor if model_card is not None else False,
     )

src/exo/shared/tests/conftest.py

@@ -7,8 +7,8 @@ import pytest
 from _pytest.logging import LogCaptureFixture
 from loguru import logger
 
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 
 
@@ -31,9 +31,8 @@ def get_pipeline_shard_metadata(
     model_id: ModelId, device_rank: int, world_size: int = 1
 ) -> ShardMetadata:
     return PipelineShardMetadata(
-        model_meta=ModelMetadata(
+        model_card=ModelCard(
             model_id=model_id,
-            pretty_name=str(model_id),
             storage_size=Memory.from_mb(100000),
             n_layers=32,
             hidden_size=1000,

src/exo/shared/tests/test_apply/test_apply_node_download.py

@@ -43,7 +43,4 @@ def test_apply_two_node_download_progress():
         NodeDownloadProgress(download_progress=event2), state
     )
 
-    # TODO: This test is failing. We should support the following:
-    # 1. Downloading multiple models concurrently on the same node (one per runner is fine).
-    # 2. Downloading a model, it completes, then downloading a different model on the same node.
     assert new_state.downloads == {NodeId("node-1"): [event1, event2]}

src/exo/shared/tests/test_state_serialization.py

@@ -1,7 +1,7 @@
 from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
-from exo.shared.types.topology import Connection
+from exo.shared.types.topology import Connection, SocketConnection
 
 
 def test_state_serialization_roundtrip() -> None:
@@ -12,9 +12,11 @@ def test_state_serialization_roundtrip() -> None:
     node_b = NodeId("node-b")
 
     connection = Connection(
-        local_node_id=node_a,
-        send_back_node_id=node_b,
-        send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
+        source=node_a,
+        sink=node_b,
+        edge=SocketConnection(
+            sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
+        ),
     )
 
     state = State()
@@ -23,5 +25,11 @@ def test_state_serialization_roundtrip() -> None:
     json_repr = state.model_dump_json()
     restored_state = State.model_validate_json(json_repr)
 
-    assert state.topology.to_snapshot() == restored_state.topology.to_snapshot()
+    assert (
+        state.topology.to_snapshot().nodes
+        == restored_state.topology.to_snapshot().nodes
+    )
+    assert set(state.topology.to_snapshot().connections) == set(
+        restored_state.topology.to_snapshot().connections
+    )
     assert restored_state.model_dump_json() == json_repr

src/exo/shared/topology.py

@@ -1,203 +1,227 @@
 import contextlib
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
 from typing import Iterable
 
 import rustworkx as rx
 from pydantic import BaseModel, ConfigDict
 
 from exo.shared.types.common import NodeId
-from exo.shared.types.profiling import ConnectionProfile, NodePerformanceProfile
-from exo.shared.types.topology import Connection, NodeInfo
+from exo.shared.types.topology import (
+    Connection,
+    Cycle,
+    RDMAConnection,
+    SocketConnection,
+)
 
 
 class TopologySnapshot(BaseModel):
-    nodes: list[NodeInfo]
-    connections: list[Connection]
+    nodes: Sequence[NodeId]
+    connections: Mapping[
+        NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]
+    ]
 
-    model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
 
+@dataclass
 class Topology:
-    def __init__(self) -> None:
-        self._graph: rx.PyDiGraph[NodeInfo, Connection] = rx.PyDiGraph()
-        self._node_id_to_rx_id_map: dict[NodeId, int] = dict()
-        self._rx_id_to_node_id_map: dict[int, NodeId] = dict()
-        self._edge_id_to_rx_id_map: dict[Connection, int] = dict()
+    _graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = field(
+        init=False, default_factory=rx.PyDiGraph
+    )
+    _vertex_indices: dict[NodeId, int] = field(init=False, default_factory=dict)
 
     def to_snapshot(self) -> TopologySnapshot:
         return TopologySnapshot(
-            nodes=list(self.list_nodes()),
-            connections=list(self.list_connections()),
+            nodes=list(self.list_nodes()), connections=self.map_connections()
         )
 
     @classmethod
     def from_snapshot(cls, snapshot: TopologySnapshot) -> "Topology":
         topology = cls()
 
-        for node in snapshot.nodes:
+        for node_id in snapshot.nodes:
             with contextlib.suppress(ValueError):
-                topology.add_node(node)
+                topology.add_node(node_id)
 
-        for connection in snapshot.connections:
-            topology.add_connection(connection)
+        for source in snapshot.connections:
+            for sink in snapshot.connections[source]:
+                for edge in snapshot.connections[source][sink]:
+                    topology.add_connection(
+                        Connection(source=source, sink=sink, edge=edge)
+                    )
 
         return topology
 
-    def add_node(self, node: NodeInfo) -> None:
-        if node.node_id in self._node_id_to_rx_id_map:
+    def add_node(self, node_id: NodeId) -> None:
+        if node_id in self._vertex_indices:
             return
-        rx_id = self._graph.add_node(node)
-        self._node_id_to_rx_id_map[node.node_id] = rx_id
-        self._rx_id_to_node_id_map[rx_id] = node.node_id
+        rx_id = self._graph.add_node(node_id)
+        self._vertex_indices[node_id] = rx_id
 
     def node_is_leaf(self, node_id: NodeId) -> bool:
         return (
-            node_id in self._node_id_to_rx_id_map
-            and len(self._graph.neighbors(self._node_id_to_rx_id_map[node_id])) == 1
+            node_id in self._vertex_indices
+            and len(self._graph.neighbors(self._vertex_indices[node_id])) <= 1
         )
 
     def neighbours(self, node_id: NodeId) -> list[NodeId]:
         return [
-            self._rx_id_to_node_id_map[rx_id]
-            for rx_id in self._graph.neighbors(self._node_id_to_rx_id_map[node_id])
+            self._graph[rx_id]
+            for rx_id in self._graph.neighbors(self._vertex_indices[node_id])
         ]
 
-    def out_edges(self, node_id: NodeId) -> list[tuple[NodeId, Connection]]:
-        if node_id not in self._node_id_to_rx_id_map:
+    def out_edges(self, node_id: NodeId) -> Iterable[Connection]:
+        if node_id not in self._vertex_indices:
             return []
-        return [
-            (self._rx_id_to_node_id_map[nid], conn)
-            for _, nid, conn in self._graph.out_edges(
-                self._node_id_to_rx_id_map[node_id]
+        return (
+            Connection(source=self._graph[source], sink=self._graph[sink], edge=edge)
+            for source, sink, edge in self._graph.out_edges(
+                self._vertex_indices[node_id]
             )
-        ]
+        )
 
     def contains_node(self, node_id: NodeId) -> bool:
-        return node_id in self._node_id_to_rx_id_map
+        return node_id in self._vertex_indices
 
-    def contains_connection(self, connection: Connection) -> bool:
-        return connection in self._edge_id_to_rx_id_map
-
-    def add_connection(
-        self,
-        connection: Connection,
-    ) -> None:
-        if connection.local_node_id not in self._node_id_to_rx_id_map:
-            self.add_node(NodeInfo(node_id=connection.local_node_id))
-        if connection.send_back_node_id not in self._node_id_to_rx_id_map:
-            self.add_node(NodeInfo(node_id=connection.send_back_node_id))
-
-        if connection in self._edge_id_to_rx_id_map:
+    def add_connection(self, conn: Connection) -> None:
+        source, sink, edge = conn.source, conn.sink, conn.edge
+        del conn
+        if edge in self.get_all_connections_between(source, sink):
             return
 
-        src_id = self._node_id_to_rx_id_map[connection.local_node_id]
-        sink_id = self._node_id_to_rx_id_map[connection.send_back_node_id]
+        if source not in self._vertex_indices:
+            self.add_node(source)
+        if sink not in self._vertex_indices:
+            self.add_node(sink)
 
-        rx_id = self._graph.add_edge(src_id, sink_id, connection)
-        self._edge_id_to_rx_id_map[connection] = rx_id
+        src_id = self._vertex_indices[source]
+        sink_id = self._vertex_indices[sink]
 
-    def list_nodes(self) -> Iterable[NodeInfo]:
-        return (self._graph[i] for i in self._graph.node_indices())
+        _ = self._graph.add_edge(src_id, sink_id, edge)
 
-    def list_connections(self) -> Iterable[Connection]:
-        return (connection for _, _, connection in self._graph.weighted_edge_list())
+    def get_all_connections_between(
+        self, source: NodeId, sink: NodeId
+    ) -> Iterable[SocketConnection | RDMAConnection]:
+        if source not in self._vertex_indices:
+            return []
+        if sink not in self._vertex_indices:
+            return []
 
-    def get_node_profile(self, node_id: NodeId) -> NodePerformanceProfile | None:
+        src_id = self._vertex_indices[source]
+        sink_id = self._vertex_indices[sink]
         try:
-            rx_idx = self._node_id_to_rx_id_map[node_id]
-            return self._graph.get_node_data(rx_idx).node_profile
-        except KeyError:
-            return None
+            return self._graph.get_all_edge_data(src_id, sink_id)
+        except rx.NoEdgeBetweenNodes:
+            return []
 
-    def update_node_profile(
-        self, node_id: NodeId, node_profile: NodePerformanceProfile
-    ) -> None:
-        rx_idx = self._node_id_to_rx_id_map[node_id]
-        self._graph[rx_idx].node_profile = node_profile
+    def list_nodes(self) -> Iterable[NodeId]:
+        return self._graph.nodes()
 
-    def update_connection_profile(self, connection: Connection) -> None:
-        rx_idx = self._edge_id_to_rx_id_map[connection]
-        self._graph.update_edge_by_index(rx_idx, connection)
+    def map_connections(
+        self,
+    ) -> Mapping[NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]]:
+        base: dict[NodeId, dict[NodeId, list[SocketConnection | RDMAConnection]]] = {}
+        for src_id, sink_id, connection in self._graph.weighted_edge_list():
+            source = self._graph[src_id]
+            sink = self._graph[sink_id]
+            if source not in base:
+                base[source] = {}
+            if sink not in base[source]:
+                base[source][sink] = []
+            base[source][sink].append(connection)
+        return base
 
-    def get_connection_profile(
-        self, connection: Connection
-    ) -> ConnectionProfile | None:
-        try:
-            rx_idx = self._edge_id_to_rx_id_map[connection]
-            return self._graph.get_edge_data_by_index(rx_idx).connection_profile
-        except KeyError:
-            return None
+    def list_connections(
+        self,
+    ) -> Iterable[Connection]:
+        return (
+            (
+                Connection(
+                    source=self._graph[src_id],
+                    sink=self._graph[sink_id],
+                    edge=connection,
+                )
+            )
+            for src_id, sink_id, connection in self._graph.weighted_edge_list()
+        )
 
     def remove_node(self, node_id: NodeId) -> None:
-        if node_id not in self._node_id_to_rx_id_map:
+        if node_id not in self._vertex_indices:
             return
 
-        for connection in self.list_connections():
-            if (
-                connection.local_node_id == node_id
-                or connection.send_back_node_id == node_id
-            ):
-                self.remove_connection(connection)
-
-        rx_idx = self._node_id_to_rx_id_map[node_id]
+        rx_idx = self._vertex_indices[node_id]
         self._graph.remove_node(rx_idx)
 
-        del self._node_id_to_rx_id_map[node_id]
-        del self._rx_id_to_node_id_map[rx_idx]
+        del self._vertex_indices[node_id]
 
-    def remove_connection(self, connection: Connection) -> None:
-        if connection not in self._edge_id_to_rx_id_map:
+    def replace_all_out_rdma_connections(
+        self, source: NodeId, new_connections: Sequence[Connection]
+    ) -> None:
+        for conn_idx in self._graph.out_edge_indices(self._vertex_indices[source]):
+            if isinstance(self._graph.get_edge_data_by_index(conn_idx), RDMAConnection):
+                self._graph.remove_edge_from_index(conn_idx)
+        for conn in new_connections:
+            self.add_connection(conn)
+
+    def remove_connection(self, conn: Connection) -> None:
+        if (
+            conn.source not in self._vertex_indices
+            or conn.sink not in self._vertex_indices
+        ):
             return
-        rx_idx = self._edge_id_to_rx_id_map[connection]
-        self._graph.remove_edge_from_index(rx_idx)
-        del self._edge_id_to_rx_id_map[connection]
+        for conn_idx in self._graph.edge_indices_from_endpoints(
+            self._vertex_indices[conn.source], self._vertex_indices[conn.sink]
+        ):
+            if self._graph.get_edge_data_by_index(conn_idx) == conn.edge:
+                self._graph.remove_edge_from_index(conn_idx)
+
+    def get_cycles(self) -> list[Cycle]:
+        """Get simple cycles in the graph, including singleton cycles"""
 
-    def get_cycles(self) -> list[list[NodeInfo]]:
         cycle_idxs = rx.simple_cycles(self._graph)
-        cycles: list[list[NodeInfo]] = []
+        cycles: list[Cycle] = []
         for cycle_idx in cycle_idxs:
-            cycle = [self._graph[idx] for idx in cycle_idx]
+            cycle = Cycle(node_ids=[self._graph[idx] for idx in cycle_idx])
             cycles.append(cycle)
-
+        for node_id in self.list_nodes():
+            cycles.append(Cycle(node_ids=[node_id]))
         return cycles
 
-    def get_cycles_tb(self) -> list[list[NodeInfo]]:
+    def get_cycles_tb(self) -> list[Cycle]:
         tb_edges = [
             (u, v, conn)
             for u, v, conn in self._graph.weighted_edge_list()
             if conn.is_thunderbolt()
         ]
 
-        tb_graph: rx.PyDiGraph[NodeInfo, Connection] = rx.PyDiGraph()
+        tb_graph: rx.PyDiGraph[NodeId, SocketConnection] = rx.PyDiGraph()
         tb_graph.add_nodes_from(self._graph.nodes())
 
         for u, v, conn in tb_edges:
-            tb_graph.add_edge(u, v, conn)
+            if isinstance(conn, SocketConnection):
+                tb_graph.add_edge(u, v, conn)
 
         cycle_idxs = rx.simple_cycles(tb_graph)
-        cycles: list[list[NodeInfo]] = []
+        cycles: list[Cycle] = []
         for cycle_idx in cycle_idxs:
-            cycle = [tb_graph[idx] for idx in cycle_idx]
+            cycle = Cycle(node_ids=[tb_graph[idx] for idx in cycle_idx])
             cycles.append(cycle)
 
         return cycles
 
-    def get_subgraph_from_nodes(self, nodes: list[NodeInfo]) -> "Topology":
-        node_idxs = [node.node_id for node in nodes]
-        rx_idxs = [self._node_id_to_rx_id_map[idx] for idx in node_idxs]
+    def get_subgraph_from_nodes(self, node_ids: list[NodeId]) -> "Topology":
         topology = Topology()
-        for rx_idx in rx_idxs:
-            topology.add_node(self._graph[rx_idx])
+        for node_id in node_ids:
+            topology.add_node(node_id)
         for connection in self.list_connections():
-            if (
-                connection.local_node_id in node_idxs
-                and connection.send_back_node_id in node_idxs
-            ):
+            if connection.source in node_ids and connection.sink in node_ids:
                 topology.add_connection(connection)
         return topology
 
-    def is_thunderbolt_cycle(self, cycle: list[NodeInfo]) -> bool:
-        node_idxs = [node.node_id for node in cycle]
-        rx_idxs = [self._node_id_to_rx_id_map[idx] for idx in node_idxs]
+    def is_thunderbolt_cycle(self, cycle: Cycle) -> bool:
+        node_idxs = [node for node in cycle]
+        rx_idxs = [self._vertex_indices[idx] for idx in node_idxs]
         for rid in rx_idxs:
             for neighbor_rid in self._graph.neighbors(rid):
                 if neighbor_rid not in rx_idxs:

src/exo/shared/types/api.py

@@ -4,17 +4,28 @@ from typing import Any, Literal
 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
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 
 FinishReason = Literal[
-    "stop", "length", "tool_calls", "content_filter", "function_call"
+    "stop", "length", "tool_calls", "content_filter", "function_call", "error"
 ]
 
 
+class ErrorInfo(BaseModel):
+    message: str
+    type: str
+    param: str | None = None
+    code: int
+
+
+class ErrorResponse(BaseModel):
+    error: ErrorInfo
+
+
 class ModelListModel(BaseModel):
     id: str
     object: str = "model"
@@ -195,7 +206,7 @@ class DeleteInstanceTaskParams(BaseModel):
 class CreateInstanceResponse(BaseModel):
     message: str
     command_id: CommandId
-    model_meta: ModelMetadata
+    model_card: ModelCard
 
 
 class DeleteInstanceResponse(BaseModel):

src/exo/shared/types/chunks.py

@@ -1,10 +1,10 @@
 from enum import Enum
 
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.api import GenerationStats
 from exo.utils.pydantic_ext import TaggedModel
 
 from .api import FinishReason
-from .models import ModelId
 
 
 class ChunkType(str, Enum):
@@ -22,6 +22,7 @@ class TokenChunk(BaseChunk):
     token_id: int
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
+    error_message: str | None = None
 
 
 class ImageChunk(BaseChunk):

src/exo/shared/types/commands.py

@@ -1,8 +1,8 @@
 from pydantic import Field
 
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.types.api import ChatCompletionTaskParams
 from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.models import ModelMetadata
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -21,7 +21,7 @@ class ChatCompletion(BaseCommand):
 
 
 class PlaceInstance(BaseCommand):
-    model_meta: ModelMetadata
+    model_card: ModelCard
     sharding: Sharding
     instance_meta: InstanceMeta
     min_nodes: int

src/exo/shared/types/common.py

@@ -16,7 +16,9 @@ class Id(str):
         cls, _source: type, handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
         # Just use a plain string schema
-        return core_schema.str_schema()
+        return core_schema.no_info_after_validator_function(
+            cls, core_schema.str_schema()
+        )
 
 
 class NodeId(Id):

src/exo/shared/types/events.py

@@ -2,14 +2,14 @@ from datetime import datetime
 
 from pydantic import Field
 
-from exo.shared.topology import Connection, NodePerformanceProfile
+from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
-from exo.shared.types.profiling import MemoryPerformanceProfile
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
+from exo.utils.info_gatherer.info_gatherer import GatheredInfo
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
@@ -76,25 +76,15 @@ class RunnerDeleted(BaseEvent):
     runner_id: RunnerId
 
 
-# TODO
-class NodeCreated(BaseEvent):
-    node_id: NodeId
-
-
 class NodeTimedOut(BaseEvent):
     node_id: NodeId
 
 
-class NodePerformanceMeasured(BaseEvent):
+# TODO: bikeshed this name
+class NodeGatheredInfo(BaseEvent):
     node_id: NodeId
     when: str  # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device
-    node_profile: NodePerformanceProfile
-
-
-class NodeMemoryMeasured(BaseEvent):
-    node_id: NodeId
-    when: str  # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device
-    memory: MemoryPerformanceProfile
+    info: GatheredInfo
 
 
 class NodeDownloadProgress(BaseEvent):
@@ -107,11 +97,11 @@ class ChunkGenerated(BaseEvent):
 
 
 class TopologyEdgeCreated(BaseEvent):
-    edge: Connection
+    conn: Connection
 
 
 class TopologyEdgeDeleted(BaseEvent):
-    edge: Connection
+    conn: Connection
 
 
 Event = (
@@ -125,10 +115,8 @@ Event = (
     | InstanceDeleted
     | RunnerStatusUpdated
     | RunnerDeleted
-    | NodeCreated
     | NodeTimedOut
-    | NodePerformanceMeasured
-    | NodeMemoryMeasured
+    | NodeGatheredInfo
     | NodeDownloadProgress
     | ChunkGenerated
     | TopologyEdgeCreated

src/exo/shared/types/models.py

@@ -1,18 +0,0 @@
-from pydantic import PositiveInt
-
-from exo.shared.types.common import Id
-from exo.shared.types.memory import Memory
-from exo.utils.pydantic_ext import CamelCaseModel
-
-
-class ModelId(Id):
-    pass
-
-
-class ModelMetadata(CamelCaseModel):
-    model_id: ModelId
-    pretty_name: str
-    storage_size: Memory
-    n_layers: PositiveInt
-    hidden_size: PositiveInt
-    supports_tensor: bool

src/exo/shared/types/multiaddr.py

@@ -1,10 +1,11 @@
 import re
 from typing import ClassVar
 
-from pydantic import BaseModel, computed_field, field_validator
+from pydantic import BaseModel, ConfigDict, computed_field, field_validator
 
 
 class Multiaddr(BaseModel):
+    model_config = ConfigDict(frozen=True)
     address: str
 
     PATTERNS: ClassVar[list[str]] = [

src/exo/shared/types/profiling.py

@@ -1,12 +1,14 @@
+from collections.abc import Sequence
 from typing import Self
 
 import psutil
 
 from exo.shared.types.memory import Memory
+from exo.shared.types.thunderbolt import ThunderboltIdentifier
 from exo.utils.pydantic_ext import CamelCaseModel
 
 
-class MemoryPerformanceProfile(CamelCaseModel):
+class MemoryUsage(CamelCaseModel):
     ram_total: Memory
     ram_available: Memory
     swap_total: Memory
@@ -44,7 +46,6 @@ class SystemPerformanceProfile(CamelCaseModel):
     sys_power: float = 0.0
     pcpu_usage: float = 0.0
     ecpu_usage: float = 0.0
-    ane_power: float = 0.0
 
 
 class NetworkInterfaceInfo(CamelCaseModel):
@@ -52,16 +53,21 @@ class NetworkInterfaceInfo(CamelCaseModel):
     ip_address: str
 
 
-class NodePerformanceProfile(CamelCaseModel):
-    model_id: str
-    chip_id: str
-    friendly_name: str
-    memory: MemoryPerformanceProfile
-    network_interfaces: list[NetworkInterfaceInfo] = []
-    system: SystemPerformanceProfile
+class NodeIdentity(CamelCaseModel):
+    """Static and slow-changing node identification data."""
+
+    model_id: str = "Unknown"
+    chip_id: str = "Unknown"
+    friendly_name: str = "Unknown"
 
 
-class ConnectionProfile(CamelCaseModel):
-    throughput: float
-    latency: float
-    jitter: float
+class NodeNetworkInfo(CamelCaseModel):
+    """Network interface information for a node."""
+
+    interfaces: Sequence[NetworkInterfaceInfo] = []
+
+
+class NodeThunderboltInfo(CamelCaseModel):
+    """Thunderbolt interface identifiers for a node."""
+
+    interfaces: Sequence[ThunderboltIdentifier] = []

src/exo/shared/types/state.py

@@ -7,7 +7,13 @@ from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
 from exo.shared.types.common import NodeId
-from exo.shared.types.profiling import NodePerformanceProfile
+from exo.shared.types.profiling import (
+    MemoryUsage,
+    NodeIdentity,
+    NodeNetworkInfo,
+    NodeThunderboltInfo,
+    SystemPerformanceProfile,
+)
 from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -35,11 +41,17 @@ class State(CamelCaseModel):
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
-    node_profiles: Mapping[NodeId, NodePerformanceProfile] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
 
+    # Granular node state mappings (update independently at different frequencies)
+    node_identities: Mapping[NodeId, NodeIdentity] = {}
+    node_memory: Mapping[NodeId, MemoryUsage] = {}
+    node_system: Mapping[NodeId, SystemPerformanceProfile] = {}
+    node_network: Mapping[NodeId, NodeNetworkInfo] = {}
+    node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
+
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()

src/exo/shared/types/thunderbolt.py

@@ -0,0 +1,81 @@
+import anyio
+from pydantic import BaseModel, Field
+
+from exo.utils.pydantic_ext import CamelCaseModel
+
+
+class ThunderboltConnection(CamelCaseModel):
+    source_uuid: str
+    sink_uuid: str
+
+
+class ThunderboltIdentifier(CamelCaseModel):
+    rdma_interface: str
+    domain_uuid: str
+
+
+## Intentionally minimal, only collecting data we care about - there's a lot more
+
+
+class _ReceptacleTag(BaseModel, extra="ignore"):
+    receptacle_id_key: str | None = None
+
+
+class _ConnectivityItem(BaseModel, extra="ignore"):
+    domain_uuid_key: str | None = None
+
+
+class ThunderboltConnectivityData(BaseModel, extra="ignore"):
+    domain_uuid_key: str | None = None
+    items: list[_ConnectivityItem] | None = Field(None, alias="_items")
+    receptacle_1_tag: _ReceptacleTag | None = None
+
+    def ident(self, ifaces: dict[str, str]) -> ThunderboltIdentifier | None:
+        if (
+            self.domain_uuid_key is None
+            or self.receptacle_1_tag is None
+            or self.receptacle_1_tag.receptacle_id_key is None
+        ):
+            return
+        tag = f"Thunderbolt {self.receptacle_1_tag.receptacle_id_key}"
+        assert tag in ifaces  # doesn't need to be an assertion but im confident
+        # if tag not in ifaces: return None
+        iface = f"rdma_{ifaces[tag]}"
+        return ThunderboltIdentifier(
+            rdma_interface=iface, domain_uuid=self.domain_uuid_key
+        )
+
+    def conn(self) -> ThunderboltConnection | None:
+        if self.domain_uuid_key is None or self.items is None:
+            return
+
+        sink_key = next(
+            (
+                item.domain_uuid_key
+                for item in self.items
+                if item.domain_uuid_key is not None
+            ),
+            None,
+        )
+        if sink_key is None:
+            return None
+
+        return ThunderboltConnection(
+            source_uuid=self.domain_uuid_key, sink_uuid=sink_key
+        )
+
+
+class ThunderboltConnectivity(BaseModel, extra="ignore"):
+    SPThunderboltDataType: list[ThunderboltConnectivityData] = []
+
+    @classmethod
+    async def gather(cls) -> list[ThunderboltConnectivityData] | None:
+        proc = await anyio.run_process(
+            ["system_profiler", "SPThunderboltDataType", "-json"], check=False
+        )
+        if proc.returncode != 0:
+            return None
+        # Saving you from PascalCase while avoiding too much pydantic
+        return ThunderboltConnectivity.model_validate_json(
+            proc.stdout
+        ).SPThunderboltDataType

src/exo/shared/types/topology.py

@@ -1,37 +1,41 @@
+from collections.abc import Iterator
+from dataclasses import dataclass
+
 from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import ConnectionProfile, NodePerformanceProfile
-from exo.utils.pydantic_ext import CamelCaseModel
+from exo.utils.pydantic_ext import FrozenModel
 
 
-class NodeInfo(CamelCaseModel):
-    node_id: NodeId
-    node_profile: NodePerformanceProfile | None = None
+@dataclass(frozen=True)
+class Cycle:
+    node_ids: list[NodeId]
+
+    def __len__(self) -> int:
+        return self.node_ids.__len__()
+
+    def __iter__(self) -> Iterator[NodeId]:
+        return self.node_ids.__iter__()
 
 
-class Connection(CamelCaseModel):
-    local_node_id: NodeId
-    send_back_node_id: NodeId
-    send_back_multiaddr: Multiaddr
-    connection_profile: ConnectionProfile | None = None
-
-    def __hash__(self) -> int:
-        return hash(
-            (
-                self.local_node_id,
-                self.send_back_node_id,
-                self.send_back_multiaddr.address,
-            )
-        )
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Connection):
-            raise ValueError("Cannot compare Connection with non-Connection")
-        return (
-            self.local_node_id == other.local_node_id
-            and self.send_back_node_id == other.send_back_node_id
-            and self.send_back_multiaddr == other.send_back_multiaddr
-        )
+class RDMAConnection(FrozenModel):
+    source_rdma_iface: str
+    sink_rdma_iface: str
 
     def is_thunderbolt(self) -> bool:
-        return str(self.send_back_multiaddr.ipv4_address).startswith("169.254")
+        return True
+
+
+class SocketConnection(FrozenModel):
+    sink_multiaddr: Multiaddr
+
+    def __hash__(self):
+        return hash(self.sink_multiaddr.ip_address)
+
+    def is_thunderbolt(self) -> bool:
+        return str(self.sink_multiaddr.ipv4_address).startswith("169.254")
+
+
+class Connection(FrozenModel):
+    source: NodeId
+    sink: NodeId
+    edge: RDMAConnection | SocketConnection

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

@@ -30,7 +30,7 @@ class MlxRingInstance(BaseInstance):
 
 
 class MlxJacclInstance(BaseInstance):
-    ibv_devices: list[list[str | None]]
+    jaccl_devices: list[list[str | None]]
     jaccl_coordinators: dict[NodeId, str]
 
 

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

@@ -1,43 +0,0 @@
-import asyncio
-from abc import ABC, abstractmethod
-from collections.abc import Coroutine
-from typing import Callable
-
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    SystemPerformanceProfile,
-)
-
-
-class ResourceCollector(ABC):
-    @abstractmethod
-    async def collect(self) -> SystemPerformanceProfile | MemoryPerformanceProfile: ...
-
-
-class SystemResourceCollector(ResourceCollector):
-    async def collect(self) -> SystemPerformanceProfile: ...
-
-
-class MemoryResourceCollector(ResourceCollector):
-    async def collect(self) -> MemoryPerformanceProfile: ...
-
-
-class ResourceMonitor:
-    data_collectors: list[ResourceCollector]
-    effect_handlers: set[
-        Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]
-    ]
-
-    async def _collect(
-        self,
-    ) -> list[SystemPerformanceProfile | MemoryPerformanceProfile]:
-        tasks: list[
-            Coroutine[None, None, SystemPerformanceProfile | MemoryPerformanceProfile]
-        ] = [collector.collect() for collector in self.data_collectors]
-        return await asyncio.gather(*tasks)
-
-    async def collect(self) -> None:
-        profiles = await self._collect()
-        for profile in profiles:
-            for effect_handler in self.effect_handlers:
-                effect_handler(profile)

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

@@ -2,8 +2,8 @@ from collections.abc import Mapping
 
 from pydantic import model_validator
 
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import Id, NodeId
-from exo.shared.types.models import ModelId
 from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
@@ -50,9 +50,7 @@ class RunnerReady(BaseRunnerStatus):
 
 
 class RunnerRunning(BaseRunnerStatus):
-    """Runner is processing requests and can accept more (continuous batching)."""
-
-    active_requests: int = 0
+    pass
 
 
 class RunnerShuttingDown(BaseRunnerStatus):

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

@@ -2,7 +2,7 @@ from enum import Enum
 
 from pydantic import Field
 
-from exo.shared.types.models import ModelMetadata
+from exo.shared.models.model_cards import ModelCard
 from exo.utils.pydantic_ext import TaggedModel
 
 
@@ -17,7 +17,7 @@ class BaseShardMetadata(TaggedModel):
     Replaces previous `Shard` object.
     """
 
-    model_meta: ModelMetadata
+    model_card: ModelCard
     device_rank: int
     world_size: int
 
@@ -41,7 +41,7 @@ class BaseShardMetadata(TaggedModel):
     def __hash__(self) -> int:
         return hash(
             (
-                self.model_meta.model_id,
+                self.model_card.model_id,
                 self.start_layer,
                 self.end_layer,
                 self.n_layers,

src/exo/utils/info_gatherer/info_gatherer.py

@@ -0,0 +1,235 @@
+import os
+import shutil
+import sys
+import tomllib
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from subprocess import CalledProcessError
+from typing import Self, cast
+
+import anyio
+from anyio import create_task_group, open_process
+from anyio.abc import TaskGroup
+from anyio.streams.buffered import BufferedByteReceiveStream
+from anyio.streams.text import TextReceiveStream
+from loguru import logger
+
+from exo.shared.constants import EXO_CONFIG_FILE
+from exo.shared.types.memory import Memory
+from exo.shared.types.profiling import (
+    MemoryUsage,
+    NetworkInterfaceInfo,
+)
+from exo.shared.types.thunderbolt import (
+    ThunderboltConnection,
+    ThunderboltConnectivity,
+    ThunderboltIdentifier,
+)
+from exo.utils.channels import Sender
+from exo.utils.pydantic_ext import TaggedModel
+
+from .macmon import MacmonMetrics
+from .system_info import get_friendly_name, get_model_and_chip, get_network_interfaces
+
+IS_DARWIN = sys.platform == "darwin"
+
+
+class StaticNodeInformation(TaggedModel):
+    """Node information that should NEVER change, to be gathered once at startup"""
+
+    model: str
+    chip: str
+
+    @classmethod
+    async def gather(cls) -> Self:
+        model, chip = await get_model_and_chip()
+        return cls(model=model, chip=chip)
+
+
+class NodeNetworkInterfaces(TaggedModel):
+    ifaces: Sequence[NetworkInterfaceInfo]
+
+
+class MacThunderboltIdentifiers(TaggedModel):
+    idents: Sequence[ThunderboltIdentifier]
+
+
+class MacThunderboltConnections(TaggedModel):
+    conns: Sequence[ThunderboltConnection]
+
+
+class NodeConfig(TaggedModel):
+    """Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there"""
+
+    @classmethod
+    async def gather(cls) -> Self | None:
+        cfg_file = anyio.Path(EXO_CONFIG_FILE)
+        await cfg_file.touch(exist_ok=True)
+        async with await cfg_file.open("rb") as f:
+            try:
+                contents = (await f.read()).decode("utf-8")
+                data = tomllib.loads(contents)
+                return cls.model_validate(data)
+            except (tomllib.TOMLDecodeError, UnicodeDecodeError):
+                logger.warning("Invalid config file, skipping...")
+                return None
+
+
+class MiscData(TaggedModel):
+    """Node information that may slowly change that doesn't fall into the other categories"""
+
+    friendly_name: str
+
+    @classmethod
+    async def gather(cls) -> Self:
+        return cls(friendly_name=await get_friendly_name())
+
+
+async def _gather_iface_map() -> dict[str, str] | None:
+    proc = await anyio.run_process(
+        ["networksetup", "-listallhardwareports"], check=False
+    )
+    if proc.returncode != 0:
+        return None
+
+    ports: dict[str, str] = {}
+    port = ""
+    for line in proc.stdout.decode("utf-8").split("\n"):
+        if line.startswith("Hardware Port:"):
+            port = line.split(": ")[1]
+        elif line.startswith("Device:"):
+            ports[port] = line.split(": ")[1]
+            port = ""
+    if "" in ports:
+        del ports[""]
+    return ports
+
+
+GatheredInfo = (
+    MacmonMetrics
+    | MemoryUsage
+    | NodeNetworkInterfaces
+    | MacThunderboltIdentifiers
+    | MacThunderboltConnections
+    | NodeConfig
+    | MiscData
+    | StaticNodeInformation
+)
+
+
+@dataclass
+class InfoGatherer:
+    info_sender: Sender[GatheredInfo]
+    interface_watcher_interval: float | None = 10
+    misc_poll_interval: float | None = 60
+    system_profiler_interval: float | None = 5 if IS_DARWIN else None
+    memory_poll_rate: float | None = None if IS_DARWIN else 1
+    macmon_interval: float | None = 1 if IS_DARWIN else None
+    _tg: TaskGroup = field(init=False, default_factory=create_task_group)
+
+    async def run(self):
+        async with self._tg as tg:
+            if IS_DARWIN:
+                if (macmon_path := shutil.which("macmon")) is not None:
+                    tg.start_soon(self._monitor_macmon, macmon_path)
+                tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
+            tg.start_soon(self._watch_system_info)
+            tg.start_soon(self._monitor_memory_usage)
+            tg.start_soon(self._monitor_misc)
+
+            nc = await NodeConfig.gather()
+            if nc is not None:
+                await self.info_sender.send(nc)
+            sni = await StaticNodeInformation.gather()
+            await self.info_sender.send(sni)
+
+    def shutdown(self):
+        self._tg.cancel_scope.cancel()
+
+    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:
+            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):
+        if self.system_profiler_interval is None:
+            return
+        iface_map = await _gather_iface_map()
+        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]
+            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))
+
+            await anyio.sleep(self.system_profiler_interval)
+
+    async def _monitor_memory_usage(self):
+        override_memory_env = os.getenv("OVERRIDE_MEMORY_MB")
+        override_memory: int | None = (
+            Memory.from_mb(int(override_memory_env)).in_bytes
+            if override_memory_env
+            else None
+        )
+        if self.memory_poll_rate is None:
+            return
+        while True:
+            await self.info_sender.send(
+                MemoryUsage.from_psutil(override_memory=override_memory)
+            )
+            await anyio.sleep(self.memory_poll_rate)
+
+    async def _watch_system_info(self):
+        if self.interface_watcher_interval is None:
+            return
+        old_nics = []
+        while True:
+            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_macmon(self, macmon_path: str):
+        if self.macmon_interval is None:
+            return
+        # macmon pipe --interval [interval in ms]
+        try:
+            async with await open_process(
+                [macmon_path, "pipe", "--interval", str(self.macmon_interval * 1000)]
+            ) as p:
+                if not p.stdout:
+                    logger.critical("MacMon closed stdout")
+                    return
+                async for text in TextReceiveStream(
+                    BufferedByteReceiveStream(p.stdout)
+                ):
+                    await self.info_sender.send(MacmonMetrics.from_raw_json(text))
+        except CalledProcessError as e:
+            stderr_msg = "no stderr"
+            stderr_output = cast(bytes | str | None, e.stderr)
+            if stderr_output is not None:
+                stderr_msg = (
+                    stderr_output.decode()
+                    if isinstance(stderr_output, bytes)
+                    else str(stderr_output)
+                )
+            logger.warning(
+                f"MacMon failed with return code {e.returncode}: {stderr_msg}"
+            )

src/exo/utils/info_gatherer/macmon.py

@@ -0,0 +1,70 @@
+from typing import Self
+
+from pydantic import BaseModel
+
+from exo.shared.types.profiling import MemoryUsage, SystemPerformanceProfile
+from exo.utils.pydantic_ext import TaggedModel
+
+
+class _TempMetrics(BaseModel, extra="ignore"):
+    """Temperature-related metrics returned by macmon."""
+
+    cpu_temp_avg: float
+    gpu_temp_avg: float
+
+
+class _MemoryMetrics(BaseModel, extra="ignore"):
+    """Memory-related metrics returned by macmon."""
+
+    ram_total: int
+    ram_usage: int
+    swap_total: int
+    swap_usage: int
+
+
+class RawMacmonMetrics(BaseModel, extra="ignore"):
+    """Complete set of metrics returned by macmon.
+
+    Unknown fields are ignored for forward-compatibility.
+    """
+
+    timestamp: str  # ignored
+    temp: _TempMetrics
+    memory: _MemoryMetrics
+    ecpu_usage: tuple[int, float]  # freq mhz, usage %
+    pcpu_usage: tuple[int, float]  # freq mhz, usage %
+    gpu_usage: tuple[int, float]  # freq mhz, usage %
+    all_power: float
+    ane_power: float
+    cpu_power: float
+    gpu_power: float
+    gpu_ram_power: float
+    ram_power: float
+    sys_power: float
+
+
+class MacmonMetrics(TaggedModel):
+    system_profile: SystemPerformanceProfile
+    memory: MemoryUsage
+
+    @classmethod
+    def from_raw(cls, raw: RawMacmonMetrics) -> Self:
+        return cls(
+            system_profile=SystemPerformanceProfile(
+                gpu_usage=raw.gpu_usage[1],
+                temp=raw.temp.gpu_temp_avg,
+                sys_power=raw.sys_power,
+                pcpu_usage=raw.pcpu_usage[1],
+                ecpu_usage=raw.ecpu_usage[1],
+            ),
+            memory=MemoryUsage.from_bytes(
+                ram_total=raw.memory.ram_total,
+                ram_available=(raw.memory.ram_total - raw.memory.ram_usage),
+                swap_total=raw.memory.swap_total,
+                swap_available=(raw.memory.swap_total - raw.memory.swap_usage),
+            ),
+        )
+
+    @classmethod
+    def from_raw_json(cls, json: str) -> Self:
+        return cls.from_raw(RawMacmonMetrics.model_validate_json(json))

src/exo/utils/info_gatherer/net_profile.py

@@ -0,0 +1,114 @@
+from collections.abc import Mapping
+
+import anyio
+import httpx
+from anyio import create_task_group
+from loguru import logger
+
+from exo.shared.topology import Topology
+from exo.shared.types.common import NodeId
+from exo.shared.types.profiling import NodeNetworkInfo
+
+REACHABILITY_ATTEMPTS = 3
+
+
+async def check_reachability(
+    target_ip: str,
+    expected_node_id: NodeId,
+    out: dict[NodeId, set[str]],
+    client: httpx.AsyncClient,
+) -> None:
+    """Check if a node is reachable at the given IP and verify its identity."""
+    if ":" in target_ip:
+        # TODO: use real IpAddress types
+        url = f"http://[{target_ip}]:52415/node_id"
+    else:
+        url = f"http://{target_ip}:52415/node_id"
+
+    remote_node_id = None
+    last_error = None
+
+    for _ in range(REACHABILITY_ATTEMPTS):
+        try:
+            r = await client.get(url)
+            if r.status_code != 200:
+                await anyio.sleep(1)
+                continue
+
+            body = r.text.strip().strip('"')
+            if not body:
+                await anyio.sleep(1)
+                continue
+
+            remote_node_id = NodeId(body)
+            break
+
+        # expected failure cases
+        except (
+            httpx.TimeoutException,
+            httpx.NetworkError,
+        ):
+            await anyio.sleep(1)
+
+        # other failures should be logged on last attempt
+        except httpx.HTTPError as e:
+            last_error = e
+            await anyio.sleep(1)
+
+    if last_error is not None:
+        logger.warning(
+            f"connect error {type(last_error).__name__} from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down"
+        )
+
+    if remote_node_id is None:
+        return
+
+    if remote_node_id != expected_node_id:
+        logger.warning(
+            f"Discovered node with unexpected node_id; "
+            f"ip={target_ip}, expected_node_id={expected_node_id}, "
+            f"remote_node_id={remote_node_id}"
+        )
+        return
+
+    if remote_node_id not in out:
+        out[remote_node_id] = set()
+    out[remote_node_id].add(target_ip)
+
+
+async def check_reachable(
+    topology: Topology,
+    self_node_id: NodeId,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
+) -> dict[NodeId, set[str]]:
+    """Check which nodes are reachable and return their IPs."""
+
+    reachable: dict[NodeId, set[str]] = {}
+
+    # these are intentionally httpx's defaults so we can tune them later
+    timeout = httpx.Timeout(timeout=5.0)
+    limits = httpx.Limits(
+        max_connections=100,
+        max_keepalive_connections=20,
+        keepalive_expiry=5,
+    )
+
+    async with (
+        httpx.AsyncClient(timeout=timeout, limits=limits) as client,
+        create_task_group() as tg,
+    ):
+        for node_id in topology.list_nodes():
+            if node_id not in node_network:
+                continue
+            if node_id == self_node_id:
+                continue
+            for iface in node_network[node_id].interfaces:
+                tg.start_soon(
+                    check_reachability,
+                    iface.ip_address,
+                    node_id,
+                    reachable,
+                    client,
+                )
+
+    return reachable

src/exo/utils/info_gatherer/tests/test_tb_parsing.py

@@ -0,0 +1,24 @@
+import sys
+
+import pytest
+
+from exo.shared.types.thunderbolt import (
+    ThunderboltConnectivity,
+)
+from exo.utils.info_gatherer.info_gatherer import (
+    _gather_iface_map,  # pyright: ignore[reportPrivateUsage]
+)
+
+
+@pytest.mark.anyio
+@pytest.mark.skipif(
+    sys.platform != "darwin", reason="Thunderbolt info can only be gathered on macos"
+)
+async def test_tb_parsing():
+    data = await ThunderboltConnectivity.gather()
+    ifaces = await _gather_iface_map()
+    assert ifaces
+    assert data
+    for datum in data:
+        datum.ident(ifaces)
+        datum.conn()

src/exo/utils/pydantic_ext.py

@@ -19,11 +19,20 @@ class CamelCaseModel(BaseModel):
         alias_generator=to_camel,
         validate_by_name=True,
         extra="forbid",
-        # I want to reenable this ASAP, but it's causing an issue with TaskStatus
         strict=True,
     )
 
 
+class FrozenModel(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        validate_by_name=True,
+        extra="forbid",
+        strict=True,
+        frozen=True,
+    )
+
+
 class TaggedModel(CamelCaseModel):
     @model_serializer(mode="wrap")
     def _serialize(self, handler: SerializerFunctionWrapHandler):

src/exo/utils/tests/test_mp_channel.py

@@ -28,9 +28,8 @@ def bar(send: MpSender[str]):
     send.close()
 
 
-# not async, just want the fail_after
 @pytest.mark.anyio
-async def test_channel_setup():
+async def test_channel_ipc():
     with fail_after(0.5):
         s, r = mp_channel[str]()
         p1 = mp.Process(target=foo, args=(r,))

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
@@ -245,12 +246,15 @@ def create_http_session(
         sock_read_timeout = 1800
         sock_connect_timeout = 60
 
-    ssl_context = ssl.create_default_context(cafile=certifi.where())
+    ssl_context = ssl.create_default_context(
+        cafile=os.getenv("SSL_CERT_FILE") or certifi.where()
+    )
     connector = aiohttp.TCPConnector(ssl=ssl_context)
 
     return aiohttp.ClientSession(
         auto_decompress=auto_decompress,
         connector=connector,
+        proxy=os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY") or None,
         timeout=aiohttp.ClientTimeout(
             total=total_timeout,
             connect=connect_timeout,
@@ -456,10 +460,10 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
     # (iii) Tensor parallel requires all files.
     return ["*"]
     try:
-        weight_map = await get_weight_map(str(shard.model_meta.model_id))
+        weight_map = await get_weight_map(str(shard.model_card.model_id))
         return get_allow_patterns(weight_map, shard)
     except Exception:
-        logger.error(f"Error getting weight map for {shard.model_meta.model_id=}")
+        logger.error(f"Error getting weight map for {shard.model_card.model_id=}")
         logger.error(traceback.format_exc())
         return ["*"]
 
@@ -522,24 +526,24 @@ 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,
 ) -> tuple[Path, RepoDownloadProgress]:
     if not skip_download:
-        logger.info(f"Downloading {shard.model_meta.model_id=}")
+        logger.info(f"Downloading {shard.model_card.model_id=}")
 
     # Handle local paths
-    if await aios.path.exists(str(shard.model_meta.model_id)):
-        logger.info(f"Using local model path {shard.model_meta.model_id}")
-        local_path = Path(str(shard.model_meta.model_id))
+    if await aios.path.exists(str(shard.model_card.model_id)):
+        logger.info(f"Using local model path {shard.model_card.model_id}")
+        local_path = Path(str(shard.model_card.model_id))
         return local_path, await download_progress_for_local_path(
-            str(shard.model_meta.model_id), shard, local_path
+            str(shard.model_card.model_id), shard, local_path
         )
 
     revision = "main"
-    target_dir = await ensure_models_dir() / str(shard.model_meta.model_id).replace(
+    target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
         "/", "--"
     )
     if not skip_download:
@@ -548,13 +552,13 @@ async def download_shard(
     if not allow_patterns:
         allow_patterns = await resolve_allow_patterns(shard)
 
-    logger.info(f"Downloading {shard.model_meta.model_id=} with {allow_patterns=}")
+    logger.info(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
 
     all_start_time = time.time()
     # TODO: currently not recursive. Some models might require subdirectories - thus this will need to be changed.
     #  Update: <- This does not seem to be the case. Yay?
     file_list = await fetch_file_list_with_cache(
-        str(shard.model_meta.model_id), revision, recursive=True
+        str(shard.model_card.model_id), revision, recursive=True
     )
     filtered_file_list = list(
         filter_repo_objects(
@@ -563,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
@@ -588,7 +592,7 @@ async def download_shard(
             else timedelta(seconds=0)
         )
         file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=str(shard.model_meta.model_id),
+            repo_id=str(shard.model_card.model_id),
             repo_revision=revision,
             file_path=file.path,
             downloaded=Memory.from_bytes(curr_bytes),
@@ -601,11 +605,11 @@ async def download_shard(
             else "in_progress",
             start_time=start_time,
         )
-        on_progress(
+        await on_progress(
             shard,
             calculate_repo_progress(
                 shard,
-                str(shard.model_meta.model_id),
+                str(shard.model_card.model_id),
                 revision,
                 file_progress,
                 all_start_time,
@@ -615,7 +619,7 @@ async def download_shard(
     for file in filtered_file_list:
         downloaded_bytes = await get_downloaded_size(target_dir / file.path)
         file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=str(shard.model_meta.model_id),
+            repo_id=str(shard.model_card.model_id),
             repo_revision=revision,
             file_path=file.path,
             downloaded=Memory.from_bytes(downloaded_bytes),
@@ -629,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),
+                str(shard.model_card.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
                 ),
             )
@@ -646,9 +657,9 @@ async def download_shard(
             *[download_with_semaphore(file) for file in filtered_file_list]
         )
     final_repo_progress = calculate_repo_progress(
-        shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
+        shard, str(shard.model_card.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:

src/exo/worker/download/impl_shard_downloader.py

@@ -1,9 +1,10 @@
 import asyncio
+from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
 from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.models.model_meta import get_model_meta
+from exo.shared.models.model_meta import get_model_card
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
@@ -19,21 +20,21 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
 
 
 async def build_base_shard(model_id: str) -> ShardMetadata:
-    model_meta = await get_model_meta(model_id)
+    model_card = await get_model_card(model_id)
     return PipelineShardMetadata(
-        model_meta=model_meta,
+        model_card=model_card,
         device_rank=0,
         world_size=1,
         start_layer=0,
-        end_layer=model_meta.n_layers,
-        n_layers=model_meta.n_layers,
+        end_layer=model_card.n_layers,
+        n_layers=model_card.n_layers,
     )
 
 
 async def build_full_shard(model_id: str) -> PipelineShardMetadata:
     base_shard = await build_base_shard(model_id)
     return PipelineShardMetadata(
-        model_meta=base_shard.model_meta,
+        model_card=base_shard.model_card,
         device_rank=base_shard.device_rank,
         world_size=base_shard.world_size,
         start_layer=base_shard.start_layer,
@@ -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,18 +85,19 @@ 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)
 
     async def ensure_shard(
         self, shard: ShardMetadata, config_only: bool = False
     ) -> Path:
-        if (shard.model_meta.model_id, shard) in self.cache:
-            return self.cache[(shard.model_meta.model_id, shard)]
+        if (shard.model_card.model_id, shard) in self.cache:
+            return self.cache[(shard.model_card.model_id, shard)]
 
         target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
-        self.cache[(shard.model_meta.model_id, shard)] = target_dir
+        self.cache[(shard.model_card.model_id, shard)] = target_dir
         return target_dir
 
     async def get_shard_download_status(
@@ -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)
 

src/exo/worker/download/shard_downloader.py

@@ -1,11 +1,12 @@
 from abc import ABC, abstractmethod
+from collections.abc import Awaitable
 from copy import copy
 from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
@@ -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
 
@@ -83,9 +86,8 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
     repo_id="noop",
     repo_revision="noop",
     shard=PipelineShardMetadata(
-        model_meta=ModelMetadata(
+        model_card=ModelCard(
             model_id=ModelId("noop"),
-            pretty_name="noope",
             storage_size=Memory.from_bytes(0),
             n_layers=1,
             hidden_size=1,

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

@@ -1,5 +1,3 @@
-from typing import Any
-
 import mlx.core as mx
 import mlx.nn as nn
 from mlx_lm.models.cache import KVCache
@@ -17,27 +15,3 @@ class Model(nn.Module):
         cache: list[KVCache] | None,
         input_embeddings: mx.array | None = None,
     ) -> mx.array: ...
-
-
-class Detokenizer:
-    def reset(self) -> None: ...
-    def add_token(self, token: int) -> None: ...
-    def finalize(self) -> None: ...
-
-    @property
-    def last_segment(self) -> str: ...
-
-
-class TokenizerWrapper:
-    bos_token: str | None
-    eos_token_ids: list[int]
-    detokenizer: Detokenizer
-
-    def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: ...
-
-    def apply_chat_template(
-        self,
-        messages_dicts: list[dict[str, Any]],
-        tokenize: bool = False,
-        add_generation_prompt: bool = True,
-    ) -> str: ...

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

@@ -1,7 +1,10 @@
+import os
+import threading
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from functools import partial
 from inspect import signature
-from typing import TYPE_CHECKING, Callable, Protocol, cast
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -10,17 +13,11 @@ from mlx.nn.layers.distributed import (
     shard_linear,
     sum_gradients,
 )
-from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
-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.gpt_oss import GptOssMoeModel
-from mlx_lm.models.gpt_oss import Model as GptOssModel
-from mlx_lm.models.llama import Model as LlamaModel
-from mlx_lm.models.minimax import Model as MiniMaxModel
-from mlx_lm.models.ministral3 import Model as Ministral3Model
 from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
 from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
@@ -29,6 +26,40 @@ from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
 from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
 
+TimeoutCallback = Callable[[], None]
+
+
+def eval_with_timeout(
+    mlx_item: Any,  # pyright: ignore[reportAny]
+    timeout_seconds: float = 60.0,
+    on_timeout: TimeoutCallback | None = None,
+) -> None:
+    """Evaluate MLX item with a hard timeout.
+
+    If on_timeout callback is provided, it will be called before terminating
+    the process. This allows the runner to send a failure event before exit.
+    """
+    completed = threading.Event()
+
+    def watchdog() -> None:
+        if not completed.wait(timeout=timeout_seconds):
+            logger.error(
+                f"mlx_item evaluation timed out after {timeout_seconds:.0f}s. "
+                "This may indicate an issue with FAST_SYNCH and tensor parallel sharding. "
+                "Terminating process."
+            )
+            if on_timeout is not None:
+                on_timeout()
+            os._exit(1)
+
+    watchdog_thread = threading.Thread(target=watchdog, daemon=True)
+    watchdog_thread.start()
+
+    try:
+        mx.eval(mlx_item)  # pyright: ignore[reportAny]
+    finally:
+        completed.set()
+
 
 class _LayerCallable(Protocol):
     """Structural type that any compatible layer must satisfy.
@@ -46,9 +77,11 @@ class CustomMlxLayer(nn.Module):
 
     def __init__(self, original_layer: _LayerCallable):
         super().__init__()
-        # Set twice to avoid __setattr__ recursion
         object.__setattr__(self, "_original_layer", original_layer)
-        self.original_layer: _LayerCallable = original_layer
+
+    @property
+    def original_layer(self) -> _LayerCallable:
+        return cast(_LayerCallable, object.__getattribute__(self, "_original_layer"))
 
     # Calls __getattr__ for any attributes not found on nn.Module (e.g. use_sliding)
     if not TYPE_CHECKING:
@@ -58,7 +91,7 @@ class CustomMlxLayer(nn.Module):
                 return super().__getattr__(name)
             except AttributeError:
                 original_layer = object.__getattribute__(self, "_original_layer")
-                return object.__getattribute__(original_layer, name)
+                return getattr(original_layer, name)
 
 
 class PipelineFirstLayer(CustomMlxLayer):
@@ -106,7 +139,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] :]
         return output
 
 
@@ -168,11 +200,21 @@ def pipeline_auto_parallel(
         inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
             start_layer:end_layer
         ]
-        inner_model_instance.swa_idx = inner_model_instance.layer_types.index(  # type: ignore
-            "sliding_attention"
+        # We can assume the model has at least one layer thanks to placement.
+        # If a layer type doesn't exist, we can set it to 0.
+        inner_model_instance.swa_idx = (
+            0
+            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
+            else inner_model_instance.layer_types.index(  # type: ignore
+                "sliding_attention"
+            )
         )
-        inner_model_instance.ga_idx = inner_model_instance.layer_types.index(  # type: ignore
-            "full_attention"
+        inner_model_instance.ga_idx = (
+            0
+            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
+            else inner_model_instance.layer_types.index(  # type: ignore
+                "full_attention"
+            )
         )
 
     _set_layers(model, layers)
@@ -181,12 +223,70 @@ def pipeline_auto_parallel(
         "Expected a list of layers after auto-parallel initialisation"
     )
 
+    return patch_pipeline_model(model, group)
+
+
+def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
+    # Patch __call__ on the model's class
+    cls = model.__class__
+    original_call = cls.__call__  # type :ignore
+    call_signature = signature(original_call)  # type :ignore
+
+    def patched_call(
+        self: T,
+        *args: object,
+        **kwargs: object,
+    ) -> mx.array:
+        logits: mx.array = original_call(self, *args, **kwargs)  # type: ignore
+        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
+            "cache", None
+        )
+
+        # Add dependency to last cache entry to ensure distributed ops are evaluated
+        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
+    return model
+
+
+def patch_tensor_model[T](model: T) -> T:
+    """Patch model's __call__ to ensure distributed ops sync during inference."""
+    cls = model.__class__
+    original_call = cls.__call__
+    call_signature = signature(original_call)
+
+    def patched_call(
+        self: T,
+        *args: object,
+        **kwargs: object,
+    ) -> mx.array:
+        logits: mx.array = original_call(self, *args, **kwargs)  # pyright: ignore[reportAny]
+        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
+            "cache", None
+        )
+
+        # Add dependency to last cache entry to ensure distributed ops are evaluated
+        if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
+            cache[-1].state = mx.depends(cache[-1].state, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
+
+        return logits
+
+    cls.__call__ = patched_call
     return model
 
 
 def tensor_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
+    timeout_seconds: float = 60.0,
+    on_timeout: TimeoutCallback | None = None,
 ) -> nn.Module:
     all_to_sharded_linear = partial(
         shard_linear,
@@ -231,37 +331,11 @@ def tensor_auto_parallel(
     if hasattr(model, "shard"):
         try:
             model.shard(group)  # type: ignore
-            return model
+            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,
-            sharded_to_all_linear,
-            all_to_sharded_linear_in_place,
-            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,
-            sharded_to_all_linear,
-            all_to_sharded_linear_in_place,
-            sharded_to_all_linear_in_place,
-        )
-    elif isinstance(model, MiniMaxModel):
-        tensor_parallel_sharding_strategy = MiniMaxShardingStrategy(
-            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)):
+    if isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -269,19 +343,13 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, GptOssModel):
-        tensor_parallel_sharding_strategy = GptOssShardingStrategy(
-            group,
-            all_to_sharded_linear,
-            sharded_to_all_linear,
-            all_to_sharded_linear_in_place,
-            sharded_to_all_linear_in_place,
-        )
-
     else:
         raise ValueError(f"Unsupported model type: {type(model)}")
 
-    return tensor_parallel_sharding_strategy.shard_model(model)
+    model = tensor_parallel_sharding_strategy.shard_model(
+        model, timeout_seconds, on_timeout
+    )
+    return patch_tensor_model(model)
 
 
 class TensorParallelShardingStrategy(ABC):
@@ -301,26 +369,12 @@ class TensorParallelShardingStrategy(ABC):
         self.N = group.size()
 
     @abstractmethod
-    def shard_model(self, model: nn.Module) -> nn.Module: ...
-
-
-class LlamaShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(self, model: nn.Module) -> nn.Module:
-        model = cast(LlamaModel, model)
-        for layer in model.layers:
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-            layer.self_attn.n_heads //= self.N
-            if layer.self_attn.n_kv_heads is not None:
-                layer.self_attn.n_kv_heads //= self.N
-
-            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)
-
-        return model
+    def shard_model(
+        self,
+        model: nn.Module,
+        timeout_seconds: float,
+        on_timeout: TimeoutCallback | None,
+    ) -> nn.Module: ...
 
 
 def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
@@ -349,93 +403,18 @@ def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
         raise ValueError("Model must have either a 'layers' or 'h' attribute")
 
 
-class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(self, model: nn.Module) -> nn.Module:
-        model = cast(DeepseekV3Model, model)
-        for layer in model.layers:
-            # Shard the self attention
-            if layer.self_attn.q_lora_rank is None:
-                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.kv_b_proj = self.all_to_sharded_linear(
-                layer.self_attn.kv_b_proj
-            )
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-            layer.self_attn.num_heads //= self.N
-
-            # Shard the MLP
-            if isinstance(layer.mlp, (DeepseekV3MLP, DeepseekV32MLP)):
-                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)
-
-            # Shard the MoE. Shard in place since the MoE should be responsible
-            # for aggregating the results.
-            else:
-                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 = ShardedDeepseekV3MoE(layer.mlp)  # type: ignore
-                layer.mlp.sharding_group = self.group
-
-        return model
-
-
-class ShardedDeepseekV3MoE(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, model: nn.Module) -> nn.Module:
-        model = cast(MiniMaxModel, model)
-        for layer in model.layers:
-            # Shard the self attention
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-            layer.self_attn.num_attention_heads //= self.N
-            layer.self_attn.num_key_value_heads //= self.N
-
-            # Shard the MoE. Shard in place since the MoE should be responsible
-            # for aggregating the results.
-            self.all_to_sharded_linear_in_place(
-                layer.block_sparse_moe.switch_mlp.gate_proj
-            )
-            self.sharded_to_all_linear_in_place(
-                layer.block_sparse_moe.switch_mlp.down_proj
-            )
-            self.all_to_sharded_linear_in_place(
-                layer.block_sparse_moe.switch_mlp.up_proj
-            )
-            layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
-            layer.block_sparse_moe.sharding_group = self.group
-
-        return model
-
-
 class QwenShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(self, model: nn.Module) -> nn.Module:
+    def shard_model(
+        self,
+        model: nn.Module,
+        timeout_seconds: float,
+        on_timeout: TimeoutCallback | None,
+    ) -> nn.Module:
         model = cast(Qwen3MoeModel, model)
         for layer in model.layers:
+            eval_with_timeout(
+                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
+            )
             # Shard the self attention
             layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -476,50 +455,3 @@ class ShardedQwenMoE(CustomMlxLayer):
         if self.sharding_group is not None:
             y = mx.distributed.all_sum(y, group=self.sharding_group)
         return y
-
-
-class GptOssShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(self, model: nn.Module) -> nn.Module:
-        model = cast(GptOssMoeModel, model)
-
-        for layer in model.layers:
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-
-            layer.self_attn.num_attention_heads //= self.N
-            layer.self_attn.num_key_value_heads //= self.N
-            layer.self_attn.num_key_value_groups = (
-                layer.self_attn.num_attention_heads
-                // layer.self_attn.num_key_value_heads
-            )
-
-            layer.self_attn.sinks = layer.self_attn.sinks[
-                layer.self_attn.num_attention_heads
-                * self.group.rank() : layer.self_attn.num_attention_heads
-                * (self.group.rank() + 1)
-            ]
-
-            self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
-            self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
-            self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
-
-            layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
-            layer.mlp.sharding_group = self.group
-
-        return model
-
-
-class ShardedGptOssMoE(CustomMlxLayer):
-    def __init__(self, layer: nn.Module):
-        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(x)
-        if self.sharding_group is not None:
-            y = mx.distributed.all_sum(y, group=self.sharding_group)
-        return y

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

@@ -1,302 +0,0 @@
-"""Batch generation engine using mlx_lm's BatchGenerator for continuous batching."""
-
-import time
-from dataclasses import dataclass, field
-
-import mlx.core as mx
-from mlx_lm.generate import BatchGenerator
-from mlx_lm.sample_utils import make_sampler
-from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
-
-from exo.shared.types.api import FinishReason, GenerationStats
-from exo.shared.types.common import CommandId
-from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams, TaskId
-from exo.shared.types.worker.runner_response import GenerationResponse
-from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import MAX_TOKENS
-from exo.worker.engines.mlx.generator.distributed_sync import share_object
-from exo.worker.engines.mlx.utils_mlx import apply_chat_template
-from exo.worker.runner.bootstrap import logger
-
-
-@dataclass
-class ActiveRequest:
-    """Tracks an active request in the batch."""
-
-    command_id: CommandId
-    task_id: TaskId
-    uid: int  # BatchGenerator's internal ID
-    detokenizer: StreamingDetokenizer
-    tokens_generated: int = 0
-    prompt_tokens: int = 0
-    start_time: float = field(default_factory=time.perf_counter)
-
-
-@dataclass
-class BatchedGenerationResponse:
-    """Response from batch engine, tagged with command_id and task_id."""
-
-    command_id: CommandId
-    task_id: TaskId
-    response: GenerationResponse
-
-
-class BatchGenerationEngine:
-    """Manages continuous batching using mlx_lm's BatchGenerator."""
-
-    def __init__(
-        self,
-        model: Model,
-        tokenizer: TokenizerWrapper,
-        group: mx.distributed.Group | None = None,
-        max_tokens: int = MAX_TOKENS,
-        completion_batch_size: int = 32,
-        prefill_batch_size: int = 8,
-        prefill_step_size: int = 2048,
-    ):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.max_tokens = max_tokens
-        self.active_requests: dict[int, ActiveRequest] = {}
-        self._pending_inserts: list[
-            tuple[CommandId, TaskId, ChatCompletionTaskParams]
-        ] = []
-        self._pending_completions: list[
-            int
-        ] = []  # UIDs completed but not yet synced/removed
-
-        self.group = group
-        self.rank = group.rank() if group else 0
-        self.is_distributed = group is not None and group.size() > 1
-
-        sampler = make_sampler(temp=0.7, top_p=1.0)
-
-        eos_tokens: set[int] = set(tokenizer.eos_token_ids or [])
-
-        self.batch_gen: BatchGenerator = BatchGenerator(
-            model=model,
-            max_tokens=max_tokens,
-            stop_tokens=eos_tokens,
-            sampler=sampler,
-            completion_batch_size=completion_batch_size,
-            prefill_batch_size=prefill_batch_size,
-            prefill_step_size=prefill_step_size,
-        )
-
-        logger.info(
-            f"BatchGenerationEngine initialized with completion_batch_size={completion_batch_size}, "
-            f"prefill_batch_size={prefill_batch_size}, distributed={self.is_distributed}"
-        )
-
-    def queue_request(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams,
-    ) -> None:
-        """Queue a request for insertion. Only rank 0 should call this.
-
-        In distributed mode, rank 0 receives tasks from the control plane and
-        queues them here. The actual insertion happens in sync_and_insert_pending()
-        which ensures all ranks insert the same requests together.
-        """
-        assert self.rank == 0, "Only rank 0 should queue requests"
-        self._pending_inserts.append((command_id, task_id, task_params))
-        logger.info(
-            f"Queued request {command_id} for insertion (pending={len(self._pending_inserts)})"
-        )
-
-    def sync_and_insert_pending(self) -> list[int]:
-        """Sync pending inserts across ranks and insert them. Returns UIDs.
-
-        This method ensures all ranks insert the same requests in the same order.
-        In non-distributed mode, it simply inserts all pending requests.
-        In distributed mode, it broadcasts pending requests from rank 0 to all ranks.
-
-        Batches all pending inserts into a single batch_gen.insert() call for
-        efficient prefill batching.
-        """
-        inserts_to_process: list[tuple[CommandId, TaskId, ChatCompletionTaskParams]]
-
-        if not self.is_distributed:
-            # Non-distributed: just insert directly from pending
-            inserts_to_process = list(self._pending_inserts)
-        else:
-            # Distributed: broadcast pending inserts from rank 0 to all ranks
-            assert self.group is not None
-            pending_data = self._pending_inserts if self.rank == 0 else None
-            synced_data = share_object(pending_data, self.rank, self.group)
-
-            if synced_data is None:
-                self._pending_inserts.clear()
-                return []
-
-            inserts_to_process = synced_data
-
-        if not inserts_to_process:
-            self._pending_inserts.clear()
-            return []
-
-        # Prepare all requests for batched insertion
-        all_tokens: list[list[int]] = []
-        all_max_tokens: list[int] = []
-        all_prompt_tokens: list[int] = []
-        request_info: list[tuple[CommandId, TaskId]] = []
-
-        for cmd_id, task_id, params in inserts_to_process:
-            prompt_str = apply_chat_template(self.tokenizer, params)
-            tokens: list[int] = self.tokenizer.encode(
-                prompt_str, add_special_tokens=False
-            )
-            max_tokens = params.max_tokens or self.max_tokens
-
-            all_tokens.append(tokens)
-            all_max_tokens.append(max_tokens)
-            all_prompt_tokens.append(len(tokens))
-            request_info.append((cmd_id, task_id))
-
-        # Single batched insert for efficient prefill
-        uids = self.batch_gen.insert(all_tokens, max_tokens=all_max_tokens)
-
-        # Track all inserted requests
-        for i, uid in enumerate(uids):
-            cmd_id, task_id = request_info[i]
-            self.active_requests[uid] = ActiveRequest(
-                command_id=cmd_id,
-                task_id=task_id,
-                uid=uid,
-                detokenizer=self.tokenizer.detokenizer,
-                prompt_tokens=all_prompt_tokens[i],
-            )
-            logger.info(
-                f"Inserted request {cmd_id} with uid={uid}, prompt_tokens={all_prompt_tokens[i]}, max_tokens={all_max_tokens[i]}"
-            )
-
-        self._pending_inserts.clear()
-        return uids
-
-    def step(self) -> list[BatchedGenerationResponse]:
-        """Run one decode step. Tracks completions but does not sync - call sync_completions() at budget boundaries."""
-        responses = self.batch_gen.next()
-        if not responses:
-            return []
-
-        results: list[BatchedGenerationResponse] = []
-
-        for r in responses:
-            uid: int = r.uid
-            req = self.active_requests.get(uid)
-            if req is None:
-                logger.warning(f"Received response for unknown uid={uid}")
-                continue
-
-            req.tokens_generated += 1
-
-            # Decode the token
-            token: int = r.token
-            req.detokenizer.add_token(token)
-            text: str = req.detokenizer.last_segment
-
-            stats: GenerationStats | None = None
-            finish_reason: FinishReason | None = None
-
-            raw_finish_reason: str | None = r.finish_reason
-            if raw_finish_reason is not None:
-                # Finalize to get remaining text
-                req.detokenizer.finalize()
-                text = req.detokenizer.last_segment
-
-                elapsed = time.perf_counter() - req.start_time
-                generation_tps = req.tokens_generated / elapsed if elapsed > 0 else 0.0
-
-                stats = GenerationStats(
-                    prompt_tps=0.0,  # Not tracked per-request in batch mode
-                    generation_tps=generation_tps,
-                    prompt_tokens=req.prompt_tokens,
-                    generation_tokens=req.tokens_generated,
-                    peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),
-                )
-
-                if raw_finish_reason == "stop":
-                    finish_reason = "stop"
-                elif raw_finish_reason == "length":
-                    finish_reason = "length"
-                else:
-                    logger.warning(f"Unknown finish_reason: {raw_finish_reason}")
-                    finish_reason = "stop"
-
-                # Track completion but don't remove yet - wait for sync_completions()
-                self._pending_completions.append(uid)
-                logger.info(
-                    f"Request {req.command_id} completed: {req.tokens_generated} tokens, {generation_tps:.2f} tps, reason={finish_reason}"
-                )
-
-            results.append(
-                BatchedGenerationResponse(
-                    command_id=req.command_id,
-                    task_id=req.task_id,
-                    response=GenerationResponse(
-                        text=text, token=token, finish_reason=finish_reason, stats=stats
-                    ),
-                )
-            )
-
-        # In non-distributed mode, clean up completions immediately
-        if not self.is_distributed:
-            self._remove_completed()
-
-        return results
-
-    def sync_completions(self) -> None:
-        """Sync and remove completed requests. Call at time budget boundaries in distributed mode."""
-        if not self.is_distributed:
-            # Non-distributed: early return if nothing to do
-            if not self._pending_completions:
-                return
-            self._remove_completed()
-            return
-
-        # Distributed mode: ALWAYS sync to ensure all ranks participate in collective op
-        # This prevents deadlock if one rank has completions and another doesn't
-        assert self.group is not None
-        synced_uids = share_object(
-            self._pending_completions if self.rank == 0 else None,
-            self.rank,
-            self.group,
-        )
-        if synced_uids:
-            self._pending_completions = synced_uids
-
-        self._remove_completed()
-
-    def _remove_completed(self) -> None:
-        """Remove completed requests from tracking."""
-        for uid in self._pending_completions:
-            if uid in self.active_requests:
-                del self.active_requests[uid]
-        self._pending_completions.clear()
-
-    @property
-    def has_active_requests(self) -> bool:
-        return bool(self.active_requests or self.batch_gen.unprocessed_prompts)
-
-    @property
-    def has_pending_inserts(self) -> bool:
-        return bool(self._pending_inserts)
-
-    @property
-    def active_count(self) -> int:
-        return len(self.active_requests)
-
-    @property
-    def pending_count(self) -> int:
-        return len(self.batch_gen.unprocessed_prompts)
-
-    @property
-    def pending_insert_count(self) -> int:
-        return len(self._pending_inserts)
-
-    @property
-    def has_pending_completions(self) -> bool:
-        return bool(self._pending_completions)

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

@@ -1,30 +0,0 @@
-"""Distributed sync utilities using mx.distributed.all_sum() to broadcast from rank 0."""
-
-# pyright: reportAny=false
-
-import pickle
-from typing import TypeVar, cast
-
-import mlx.core as mx
-
-T = TypeVar("T")
-
-
-def share_object(obj: T | None, rank: int, group: mx.distributed.Group) -> T | None:
-    """Broadcast object from rank 0 to all ranks. Two-phase: size then data."""
-    if rank == 0:
-        if obj is None:
-            mx.eval(mx.distributed.all_sum(mx.array([0]), group=group))
-            return None
-        data = mx.array(list(pickle.dumps(obj)), dtype=mx.uint8)
-        mx.eval(mx.distributed.all_sum(mx.array([data.size]), group=group))
-        mx.eval(mx.distributed.all_sum(data, group=group))
-        return obj
-    else:
-        size = int(mx.distributed.all_sum(mx.array([0]), group=group).item())
-        if size == 0:
-            return None
-        data = mx.zeros(size, dtype=mx.uint8)
-        data = mx.distributed.all_sum(data, group=group)
-        mx.eval(data)
-        return cast(T, pickle.loads(bytes(cast(list[int], data.tolist()))))

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

@@ -119,6 +119,7 @@ def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
     task: ChatCompletionTaskParams,
+    prompt: str,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -130,11 +131,6 @@ def mlx_generate(
     if task.seed is not None:
         mx.random.seed(task.seed)
 
-    prompt = apply_chat_template(
-        tokenizer=tokenizer,
-        chat_task_data=task,
-    )
-
     caches = make_kv_cache(model=model)
 
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []

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

@@ -1,104 +0,0 @@
-"""Time budget iterator for controlling generation loop timing in distributed mode.
-
-Based on mlx-lm's TimeBudget pattern - runs for a time budget then syncs,
-rather than syncing every token. This reduces distributed sync overhead.
-"""
-
-import time
-from typing import Iterator
-
-import mlx.core as mx
-
-from exo.worker.runner.bootstrap import logger
-
-generation_stream = mx.new_stream(mx.default_device())
-
-
-class TimeBudget(Iterator[None]):
-    """Controls generation loop timing, syncing across ranks periodically.
-
-    In distributed mode, periodically syncs timing across all ranks to
-    dynamically adjust iteration count based on actual performance.
-
-    In non-distributed mode, simply runs for the time budget.
-
-    Usage:
-        for _ in TimeBudget(budget=0.5):
-            batch_engine.step()
-            # ... process responses ...
-    """
-
-    def __init__(
-        self,
-        budget: float = 0.5,
-        iterations: int = 25,
-        sync_frequency: int = 10,
-        group: mx.distributed.Group | None = None,
-    ):
-        """Initialize TimeBudget.
-
-        Args:
-            budget: Time budget in seconds before yielding control
-            iterations: Initial number of iterations per budget period (distributed only)
-            sync_frequency: How often to sync timing across ranks (distributed only)
-            group: Distributed group, or None for non-distributed mode
-        """
-        self._budget = budget
-        self._iterations = iterations
-        self._sync_frequency = sync_frequency
-        self._group = group
-        self._is_distributed = group is not None and group.size() > 1
-
-        # Runtime state
-        self._start: float = 0.0
-        self._current_iterations: int = 0
-        self._loops: int = 0
-        self._time_spent: float = 0.0
-
-    def __iter__(self) -> "TimeBudget":
-        self._start = time.perf_counter()
-        self._current_iterations = 0
-        return self
-
-    def __next__(self) -> None:
-        if not self._is_distributed:
-            # Non-distributed: just check time budget
-            if time.perf_counter() - self._start > self._budget:
-                raise StopIteration()
-            return None
-
-        # Distributed mode: iteration-based with periodic timing sync
-        self._current_iterations += 1
-        if self._current_iterations > self._iterations:
-            self._loops += 1
-            self._time_spent += time.perf_counter() - self._start
-
-            if self._loops % self._sync_frequency == 0:
-                # Sync timing across all ranks
-                assert self._group is not None
-                with mx.stream(generation_stream):
-                    time_array = mx.array([self._time_spent], dtype=mx.float32)
-                    total_time = mx.distributed.all_sum(time_array, group=self._group)
-                    mx.eval(total_time)
-                    loop_time = float(total_time.item())
-
-                avg_loop_time = loop_time / (self._group.size() * self._sync_frequency)
-
-                if avg_loop_time > 0:
-                    factor = self._budget / avg_loop_time
-                    self._iterations = max(round(self._iterations * factor), 1)
-                    logger.debug(
-                        f"TimeBudget adjusted iterations to {self._iterations}"
-                    )
-
-                self._loops = 0
-                self._time_spent = 0.0
-
-            raise StopIteration()
-
-        return None
-
-    @property
-    def iterations(self) -> int:
-        """Current iterations per budget period."""
-        return self._iterations

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

@@ -57,6 +57,8 @@ from exo.shared.types.worker.shards import (
 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,
+    eval_with_timeout,
     pipeline_auto_parallel,
     tensor_auto_parallel,
 )
@@ -73,7 +75,7 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     return Memory.from_float_kb(
         (model_shard_meta.end_layer - model_shard_meta.start_layer)
         / model_shard_meta.n_layers
-        * model_shard_meta.model_meta.storage_size.in_kb
+        * model_shard_meta.model_card.storage_size.in_kb
         / (
             1
             if isinstance(model_shard_meta, PipelineShardMetadata)
@@ -82,6 +84,10 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     )
 
 
+class ModelLoadingTimeoutError(Exception):
+    pass
+
+
 def mx_barrier(group: Group | None = None):
     mx.eval(
         mx.distributed.all_sum(
@@ -145,22 +151,28 @@ def mlx_distributed_init(
                 group = mx.distributed.init(backend="ring", strict=True)
 
             case MlxJacclInstance(
-                ibv_devices=ibv_devices, jaccl_coordinators=jaccl_coordinators
+                jaccl_devices=jaccl_devices, jaccl_coordinators=jaccl_coordinators
             ):
+                assert all(
+                    jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
+                )
                 # Use RDMA connectivity matrix
                 coordination_file = (
                     f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
                 )
-                ibv_devices_json = json.dumps(ibv_devices)
+                jaccl_devices_json = json.dumps(jaccl_devices)
 
                 with open(coordination_file, "w") as f:
-                    _ = f.write(ibv_devices_json)
+                    _ = f.write(jaccl_devices_json)
 
                 jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
 
-                logger.info(f"rank {rank} MLX_IBV_DEVICES: {ibv_devices_json}")
+                # TODO: update once upstream fixes
+                logger.info(
+                    f"rank {rank} MLX_JACCL_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
+                )
                 logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
-                os.environ["MLX_IBV_DEVICES"] = coordination_file
+                os.environ["MLX_JACCL_DEVICES"] = coordination_file
                 os.environ["MLX_RANK"] = str(rank)
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
@@ -188,11 +200,13 @@ def initialize_mlx(
 
 
 def load_mlx_items(
-    bound_instance: BoundInstance, group: Group | None
+    bound_instance: BoundInstance,
+    group: Group | None,
+    on_timeout: TimeoutCallback | None = None,
 ) -> tuple[Model, TokenizerWrapper]:
     if group is None:
         logger.info(f"Single device used for {bound_instance.instance}")
-        model_path = build_model_path(bound_instance.bound_shard.model_meta.model_id)
+        model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
         start_time = time.perf_counter()
         model, _ = load_model(model_path, strict=True)
         end_time = time.perf_counter()
@@ -202,7 +216,9 @@ def load_mlx_items(
     else:
         logger.info("Starting distributed init")
         start_time = time.perf_counter()
-        model, tokenizer = shard_and_load(bound_instance.bound_shard, group=group)
+        model, tokenizer = shard_and_load(
+            bound_instance.bound_shard, group=group, on_timeout=on_timeout
+        )
         end_time = time.perf_counter()
         logger.info(
             f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
@@ -216,8 +232,9 @@ def load_mlx_items(
 def shard_and_load(
     shard_metadata: ShardMetadata,
     group: Group,
+    on_timeout: TimeoutCallback | None = None,
 ) -> tuple[nn.Module, TokenizerWrapper]:
-    model_path = build_model_path(shard_metadata.model_meta.model_id)
+    model_path = build_model_path(shard_metadata.model_card.model_id)
 
     model, _ = load_model(model_path, lazy=True, strict=False)
     logger.debug(model)
@@ -244,15 +261,23 @@ def shard_and_load(
 
     logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")
 
+    # Estimate timeout based on model size
+    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "60"))
+    model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3)
+    timeout_seconds = base_timeout + model_size_gb / 5
+    logger.info(
+        f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
+        f"(model size: {model_size_gb:.1f}GB)"
+    )
+
     match shard_metadata:
         case TensorShardMetadata():
             logger.info(f"loading model from {model_path} with tensor parallelism")
-            model = tensor_auto_parallel(model, group)
+            model = tensor_auto_parallel(model, group, timeout_seconds, on_timeout)
         case PipelineShardMetadata():
             logger.info(f"loading model from {model_path} with pipeline parallelism")
             model = pipeline_auto_parallel(model, group, shard_metadata)
-
-    mx.eval(model.parameters())
+            eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
 
     # TODO: Do we need this?
     mx.eval(model)
@@ -268,7 +293,7 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
-    return load_tokenizer_for_model_id(shard_metadata.model_meta.model_id, model_path)
+    return load_tokenizer_for_model_id(shard_metadata.model_card.model_id, model_path)
 
 
 def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
@@ -287,6 +312,9 @@ def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
     model_id_lower = model_id.lower()
     if "kimi-k2" in model_id_lower:
         return [163586]
+    elif "glm-4.7-flash" in model_id_lower:
+        # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|>
+        return [154820, 154827, 154829]
     elif "glm" in model_id_lower:
         return [151336, 151329, 151338]
     return None
@@ -371,6 +399,16 @@ def apply_chat_template(
     return prompt
 
 
+def detect_thinking_prompt_suffix(prompt: str, tokenizer: TokenizerWrapper) -> bool:
+    """
+    Detect if prompt ends with a thinking opening tag that should be
+    prepended to the output stream.
+    """
+    think_token = tokenizer.think_start
+
+    return think_token is not None and prompt.rstrip().endswith(think_token)
+
+
 class NullKVCache(KVCache):
     """
     A KVCache that pretends to exist but holds zero tokens.

src/exo/worker/main.py

@@ -8,6 +8,7 @@ from loguru import logger
 
 from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
 from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import ForwarderCommand, RequestEventLog
 from exo.shared.types.common import NodeId, SessionId
 from exo.shared.types.events import (
@@ -16,16 +17,13 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     NodeDownloadProgress,
-    NodeMemoryMeasured,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
     TaskCreated,
     TaskStatusUpdated,
     TopologyEdgeCreated,
     TopologyEdgeDeleted,
 )
-from exo.shared.types.models import ModelId
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import MemoryPerformanceProfile, NodePerformanceProfile
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     CreateRunner,
@@ -34,7 +32,7 @@ from exo.shared.types.tasks import (
     Task,
     TaskStatus,
 )
-from exo.shared.types.topology import Connection
+from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
     DownloadOngoing,
@@ -45,14 +43,14 @@ 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.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
-from exo.worker.utils import start_polling_memory_metrics, start_polling_node_metrics
-from exo.worker.utils.net_profile import check_reachable
 
 
 class Worker:
@@ -86,7 +84,7 @@ class Worker:
         self.state: State = State()
         self.download_status: dict[ModelId, DownloadProgress] = {}
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
-        self._tg: TaskGroup | None = None
+        self._tg: TaskGroup = create_task_group()
 
         self._nack_cancel_scope: CancelScope | None = None
         self._nack_attempts: int = 0
@@ -98,37 +96,13 @@ class Worker:
     async def run(self):
         logger.info("Starting Worker")
 
-        # TODO: CLEANUP HEADER
-        async def resource_monitor_callback(
-            node_performance_profile: NodePerformanceProfile,
-        ) -> None:
-            await self.event_sender.send(
-                NodePerformanceMeasured(
-                    node_id=self.node_id,
-                    node_profile=node_performance_profile,
-                    when=str(datetime.now(tz=timezone.utc)),
-                ),
-            )
+        info_send, info_recv = channel[GatheredInfo]()
+        info_gatherer: InfoGatherer = InfoGatherer(info_send)
 
-        async def memory_monitor_callback(
-            memory_profile: MemoryPerformanceProfile,
-        ) -> None:
-            await self.event_sender.send(
-                NodeMemoryMeasured(
-                    node_id=self.node_id,
-                    memory=memory_profile,
-                    when=str(datetime.now(tz=timezone.utc)),
-                )
-            )
-
-        # END CLEANUP
-
-        async with create_task_group() as tg:
-            self._tg = tg
+        async with self._tg as tg:
+            tg.start_soon(info_gatherer.run)
+            tg.start_soon(self._forward_info, info_recv)
             tg.start_soon(self.plan_step)
-            tg.start_soon(start_polling_node_metrics, resource_monitor_callback)
-
-            tg.start_soon(start_polling_memory_metrics, memory_monitor_callback)
             tg.start_soon(self._emit_existing_download_progress)
             tg.start_soon(self._connection_message_event_writer)
             tg.start_soon(self._resend_out_for_delivery)
@@ -142,6 +116,17 @@ class Worker:
         for runner in self.runners.values():
             runner.shutdown()
 
+    async def _forward_info(self, recv: Receiver[GatheredInfo]):
+        with recv as info_stream:
+            async for info in info_stream:
+                await self.event_sender.send(
+                    NodeGatheredInfo(
+                        node_id=self.node_id,
+                        when=str(datetime.now(tz=timezone.utc)),
+                        info=info,
+                    )
+                )
+
     async def _event_applier(self):
         with self.global_event_receiver as events:
             async for f_event in events:
@@ -161,7 +146,6 @@ class Worker:
                     self._nack_cancel_scope is None
                     or self._nack_cancel_scope.cancel_called
                 ):
-                    assert self._tg
                     # Request the next index.
                     self._tg.start_soon(
                         self._nack_request, self.state.last_event_applied_idx + 1
@@ -202,11 +186,11 @@ class Worker:
                         )
                     )
                 case DownloadModel(shard_metadata=shard):
-                    if shard.model_meta.model_id not in self.download_status:
+                    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_meta.model_id] = progress
+                        self.download_status[shard.model_card.model_id] = progress
                         await self.event_sender.send(
                             NodeDownloadProgress(download_progress=progress)
                         )
@@ -221,7 +205,7 @@ class Worker:
                             node_id=self.node_id,
                             total_bytes=initial_progress.total_bytes,
                         )
-                        self.download_status[shard.model_meta.model_id] = progress
+                        self.download_status[shard.model_card.model_id] = progress
                         await self.event_sender.send(
                             NodeDownloadProgress(download_progress=progress)
                         )
@@ -252,8 +236,7 @@ class Worker:
                     await self.runners[self._task_to_runner_id(task)].start_task(task)
 
     def shutdown(self):
-        if self._tg:
-            self._tg.cancel_scope.cancel()
+        self._tg.cancel_scope.cancel()
 
     def _task_to_runner_id(self, task: Task):
         instance = self.state.instances[task.instance_id]
@@ -270,24 +253,28 @@ class Worker:
         match msg.connection_type:
             case ConnectionMessageType.Connected:
                 return TopologyEdgeCreated(
-                    edge=Connection(
-                        local_node_id=self.node_id,
-                        send_back_node_id=msg.node_id,
-                        send_back_multiaddr=Multiaddr(
-                            address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
                         ),
-                    )
+                    ),
                 )
 
             case ConnectionMessageType.Disconnected:
                 return TopologyEdgeDeleted(
-                    edge=Connection(
-                        local_node_id=self.node_id,
-                        send_back_node_id=msg.node_id,
-                        send_back_multiaddr=Multiaddr(
-                            address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
                         ),
-                    )
+                    ),
                 )
 
     async def _nack_request(self, since_idx: int) -> None:
@@ -336,7 +323,6 @@ class Worker:
             event_sender=self.event_sender.clone(),
         )
         self.runners[task.bound_instance.bound_runner_id] = runner
-        assert self._tg
         self._tg.start_soon(runner.run)
         return runner
 
@@ -353,14 +339,13 @@ class Worker:
                 initial_progress
             ),
         )
-        self.download_status[task.shard_metadata.model_meta.model_id] = status
+        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
 
-        # TODO: i hate callbacks
-        def download_progress_callback(
+        async def download_progress_callback(
             shard: ShardMetadata, progress: RepoDownloadProgress
         ) -> None:
             nonlocal self
@@ -371,12 +356,11 @@ class Worker:
                     node_id=self.node_id,
                     total_bytes=progress.total_bytes,
                 )
-                self.download_status[shard.model_meta.model_id] = status
-                # Footgun!
-                self.event_sender.send_nowait(
+                self.download_status[shard.model_card.model_id] = status
+                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
                     )
@@ -392,14 +376,13 @@ class Worker:
                         progress
                     ),
                 )
-                self.download_status[shard.model_meta.model_id] = status
-                self.event_sender.send_nowait(
+                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)
-        assert self._tg
         self._tg.start_soon(self.shard_downloader.ensure_shard, task.shard_metadata)
 
     async def _forward_events(self) -> None:
@@ -420,36 +403,46 @@ class Worker:
 
     async def _poll_connection_updates(self):
         while True:
-            # TODO: EdgeDeleted
-            edges = set(self.state.topology.list_connections())
-            conns = await check_reachable(self.state.topology, self.node_id)
+            edges = set(
+                conn.edge for conn in self.state.topology.out_edges(self.node_id)
+            )
+            conns = await check_reachable(
+                self.state.topology,
+                self.node_id,
+                self.state.node_network,
+            )
             for nid in conns:
                 for ip in conns[nid]:
-                    if "127.0.0.1" in ip or "localhost" in ip:
-                        logger.warning(
-                            f"Loopback connection should not happen: {ip=} for {nid=}"
-                        )
-
-                    edge = Connection(
-                        local_node_id=self.node_id,
-                        send_back_node_id=nid,
+                    edge = SocketConnection(
                         # nonsense multiaddr
-                        send_back_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
+                        sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
                         if "." in ip
                         # nonsense multiaddr
                         else Multiaddr(address=f"/ip6/{ip}/tcp/52415"),
                     )
                     if edge not in edges:
                         logger.debug(f"ping discovered {edge=}")
-                        await self.event_sender.send(TopologyEdgeCreated(edge=edge))
+                        await self.event_sender.send(
+                            TopologyEdgeCreated(
+                                conn=Connection(
+                                    source=self.node_id, sink=nid, edge=edge
+                                )
+                            )
+                        )
 
-            for nid, conn in self.state.topology.out_edges(self.node_id):
+            for conn in self.state.topology.out_edges(self.node_id):
+                if not isinstance(conn.edge, SocketConnection):
+                    continue
+                # ignore mDNS discovered connections
+                if conn.edge.sink_multiaddr.port != 52415:
+                    continue
                 if (
-                    nid not in conns
-                    or conn.send_back_multiaddr.ip_address not in conns.get(nid, set())
+                    conn.sink not in conns
+                    or conn.edge.sink_multiaddr.ip_address
+                    not in conns.get(conn.sink, set())
                 ):
                     logger.debug(f"ping failed to discover {conn=}")
-                    await self.event_sender.send(TopologyEdgeDeleted(edge=conn))
+                    await self.event_sender.send(TopologyEdgeDeleted(conn=conn))
 
             await anyio.sleep(10)
 
@@ -483,7 +476,7 @@ class Worker:
                     else:
                         continue
 
-                    self.download_status[progress.shard.model_meta.model_id] = status
+                    self.download_status[progress.shard.model_card.model_id] = status
                     await self.event_sender.send(
                         NodeDownloadProgress(download_progress=status)
                     )

src/exo/worker/plan.py

@@ -2,8 +2,8 @@
 
 from collections.abc import Mapping, Sequence
 
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import NodeId
-from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -114,7 +114,7 @@ def _model_needs_download(
     download_status: Mapping[ModelId, DownloadProgress],
 ) -> DownloadModel | None:
     for runner in runners.values():
-        model_id = runner.bound_instance.bound_shard.model_meta.model_id
+        model_id = runner.bound_instance.bound_shard.model_card.model_id
         if isinstance(runner.status, RunnerIdle) and (
             model_id not in download_status
             or not isinstance(
@@ -191,7 +191,7 @@ def _load_model(
             nid in global_download_status
             and any(
                 isinstance(dp, DownloadCompleted)
-                and dp.shard_metadata.model_meta.model_id == shard_assignments.model_id
+                and dp.shard_metadata.model_card.model_id == shard_assignments.model_id
                 for dp in global_download_status[nid]
             )
             for nid in shard_assignments.node_to_runner
@@ -277,14 +277,12 @@ def _pending_tasks(
             # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough
             # however, realistically the task status should be set to completed by the LAST runner, so this is a true race
             # the actual solution is somewhat deeper than this bypass - TODO!
-            # Also skip tasks in pending to prevent duplicate forwarding with continuous batching
-            if task.task_id in runner.completed or task.task_id in runner.pending:
+            if task.task_id in runner.completed:
                 continue
 
             # TODO: Check ordering aligns with MLX distributeds expectations.
 
-            # Allow forwarding tasks when runner is Ready or Running (for continuous batching)
-            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
+            if isinstance(runner.status, RunnerReady) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/bootstrap.py

@@ -17,15 +17,23 @@ def entrypoint(
     task_receiver: MpReceiver[Task],
     _logger: "loguru.Logger",
 ) -> None:
-    if (
-        isinstance(bound_instance.instance, MlxJacclInstance)
-        and len(bound_instance.instance.ibv_devices) >= 2
+    fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
+    if fast_synch_override == "on" or (
+        fast_synch_override != "off"
+        and (
+            isinstance(bound_instance.instance, MlxJacclInstance)
+            and len(bound_instance.instance.jaccl_devices) >= 2
+        )
     ):
         os.environ["MLX_METAL_FAST_SYNCH"] = "1"
+    else:
+        os.environ["MLX_METAL_FAST_SYNCH"] = "0"
 
     global logger
     logger = _logger
 
+    logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
+
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         from exo.worker.runner.runner import main

src/exo/worker/runner/runner.py

@@ -1,8 +1,16 @@
-import gc
 import time
+from collections.abc import Generator
+from functools import cache
 
 import mlx.core as mx
-from anyio import WouldBlock
+from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+    HarmonyEncodingName,
+    Role,
+    StreamableParser,
+    load_harmony_encoding,
+)
 
 from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.chunks import TokenChunk
@@ -23,6 +31,9 @@ from exo.shared.types.tasks import (
     TaskStatus,
 )
 from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import (
+    GenerationResponse,
+)
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
@@ -38,10 +49,10 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx.generator.batch_engine import BatchGenerationEngine
-from exo.worker.engines.mlx.generator.generate import warmup_inference
-from exo.worker.engines.mlx.generator.time_budget import TimeBudget
+from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    detect_thinking_prompt_suffix,
     initialize_mlx,
     load_mlx_items,
     mlx_force_oom,
@@ -59,6 +70,7 @@ def main(
         bound_instance.bound_runner_id,
         bound_instance.bound_shard,
     )
+    device_rank = shard_metadata.device_rank
     logger.info("hello from the runner")
     if getattr(shard_metadata, "immediate_exception", False):
         raise Exception("Fake exception - runner failed to spin up.")
@@ -70,318 +82,251 @@ def main(
     model = None
     tokenizer = None
     group = None
-    batch_engine: BatchGenerationEngine | None = None
-    pending_shutdown: Shutdown | None = None
 
     current_status: RunnerStatus = RunnerIdle()
-
-    def send_status(status: RunnerStatus) -> None:
-        event_sender.send(
-            RunnerStatusUpdated(runner_id=runner_id, runner_status=status)
-        )
-
     logger.info("runner created")
-    send_status(current_status)
-
-    def handle_task(task: Task, is_deferred: bool = False) -> bool:
-        nonlocal current_status, model, tokenizer, group, batch_engine, pending_shutdown
-
-        # For Shutdown, check if we need to defer BEFORE sending Running/Acknowledged
-        if (
-            isinstance(task, Shutdown)
-            and not is_deferred
-            and batch_engine is not None
-            and (batch_engine.has_active_requests or batch_engine.has_pending_inserts)
-        ):
-            logger.info("deferring shutdown until active requests complete")
-            pending_shutdown = task
-            return True
-
-        event_sender.send(
-            TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
-        )
-        event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-        match task:
-            case ConnectToGroup() if isinstance(
-                current_status, (RunnerIdle, RunnerFailed)
-            ):
-                logger.info("runner connecting")
-                current_status = RunnerConnecting()
-                send_status(current_status)
-                group = initialize_mlx(bound_instance)
-
-                logger.info("runner connected")
-                current_status = RunnerConnected()
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
-                send_status(current_status)
-
-            case LoadModel() if (
-                isinstance(current_status, RunnerConnected) and group is not None
-            ) or (isinstance(current_status, RunnerIdle) and group is None):
-                current_status = RunnerLoading()
-                logger.info("runner loading")
-                send_status(current_status)
-
-                model, tokenizer = load_mlx_items(bound_instance, group)
-
-                current_status = RunnerLoaded()
-                logger.info("runner loaded")
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
-                send_status(current_status)
-
-            case StartWarmup() if isinstance(current_status, RunnerLoaded):
-                assert model is not None
-                assert tokenizer is not None
-                current_status = RunnerWarmingUp()
-                logger.info("runner warming up")
-                send_status(current_status)
-
-                logger.info(f"warming up inference for instance: {instance}")
-                toks = warmup_inference(model=model, tokenizer=tokenizer)
-                logger.info(f"warmed up by generating {toks} tokens")
-                logger.info(
-                    f"runner initialized in {time.time() - setup_start_time} seconds"
-                )
-
-                batch_engine = BatchGenerationEngine(
-                    model=model, tokenizer=tokenizer, group=group
-                )
-
-                current_status = RunnerReady()
-                logger.info("runner ready")
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
-                send_status(current_status)
-
-            case ChatCompletion(task_params=task_params, command_id=command_id) if (
-                isinstance(current_status, (RunnerReady, RunnerRunning))
-            ):
-                assert batch_engine is not None
-
-                # In distributed mode, only rank 0 should queue requests
-                # Other ranks should skip - they'll participate in sync_and_insert_pending()
-                is_distributed_mode = group is not None and group.size() > 1
-                if is_distributed_mode and shard_metadata.device_rank != 0:
-                    logger.debug(
-                        f"Rank {shard_metadata.device_rank} skipping ChatCompletionTask (only rank 0 queues)"
-                    )
-                    return True
-
-                if task_params.messages and task_params.messages[0].content is not None:
-                    _check_for_debug_prompts(task_params.messages[0].content)
-
-                # Queue the request - actual insertion happens in sync_and_insert_pending()
-                batch_engine.queue_request(
-                    command_id=command_id, task_id=task.task_id, task_params=task_params
-                )
-
-                # Status will be updated after actual insertion in the main loop
-                # For now, set to RunnerRunning to indicate we're processing
-                current_status = RunnerRunning(
-                    active_requests=batch_engine.active_count
-                    + batch_engine.pending_insert_count
-                )
-                send_status(current_status)
-
-            case Shutdown():
-                current_status = RunnerShuttingDown()
-                logger.info("runner shutting down")
-                send_status(current_status)
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
-                current_status = RunnerShutdown()
-                send_status(current_status)
-                return False
-
-            case _:
-                raise ValueError(
-                    f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
-                )
-
-        return True
-
+    event_sender.send(
+        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+    )
     with task_receiver as tasks:
-        running = True
-        is_rank_0 = shard_metadata.device_rank == 0
-
-        while running:
-            # Use batch_engine.is_distributed since it's set correctly after group initialization
-            # (the group variable is None at loop start, but set by ConnectToGroup task)
-            if batch_engine is not None and batch_engine.is_distributed:
-                assert group is not None
-                assert batch_engine is not None
-
-                # Distributed mode: tasks wake up all ranks, then we sync and generate
-
-                # Check deferred shutdown FIRST - all ranks must check and process together
-                # This must run before any collective operations to prevent deadlock
-                if (
-                    pending_shutdown is not None
-                    and not batch_engine.has_active_requests
-                    and not batch_engine.has_pending_inserts
+        for task in tasks:
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
+            )
+            event_sender.send(TaskAcknowledged(task_id=task.task_id))
+            match task:
+                case ConnectToGroup() if isinstance(
+                    current_status, (RunnerIdle, RunnerFailed)
                 ):
-                    handle_task(pending_shutdown, is_deferred=True)
-                    running = False
-                    continue
+                    logger.info("runner connecting")
+                    current_status = RunnerConnecting()
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+                    group = initialize_mlx(bound_instance)
 
-                # When idle, block waiting for task (exo sends tasks to all ranks)
-                # When active, poll non-blocking to batch incoming requests
-                if (
-                    not batch_engine.has_active_requests
-                    and not batch_engine.has_pending_inserts
+                    logger.info("runner connected")
+                    current_status = RunnerConnected()
+
+                # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
+                case LoadModel() if (
+                    isinstance(current_status, RunnerConnected) and group is not None
+                ) or (isinstance(current_status, RunnerIdle) and group is None):
+                    current_status = RunnerLoading()
+                    logger.info("runner loading")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+
+                    def on_model_load_timeout() -> None:
+                        event_sender.send(
+                            RunnerStatusUpdated(
+                                runner_id=runner_id,
+                                runner_status=RunnerFailed(
+                                    error_message="Model loading timed out"
+                                ),
+                            )
+                        )
+                        time.sleep(0.5)
+
+                    model, tokenizer = load_mlx_items(
+                        bound_instance, group, on_timeout=on_model_load_timeout
+                    )
+
+                    current_status = RunnerLoaded()
+                    logger.info("runner loaded")
+                case StartWarmup() if isinstance(current_status, RunnerLoaded):
+                    assert model
+                    assert tokenizer
+                    current_status = RunnerWarmingUp()
+                    logger.info("runner warming up")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+
+                    logger.info(f"warming up inference for instance: {instance}")
+                    toks = warmup_inference(
+                        model=model,
+                        tokenizer=tokenizer,
+                        # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
+                    )
+                    logger.info(f"warmed up by generating {toks} tokens")
+                    logger.info(
+                        f"runner initialized in {time.time() - setup_start_time} seconds"
+                    )
+                    current_status = RunnerReady()
+                    logger.info("runner ready")
+                case ChatCompletion(task_params=task_params, command_id=command_id) if (
+                    isinstance(current_status, RunnerReady)
                 ):
-                    # IDLE: Block until task arrives (all ranks receive the same task)
-                    task = tasks.receive()
-                    task_result = handle_task(task)
-                    if not task_result:
-                        running = False
-                        continue
-                else:
-                    # ACTIVE: Poll for new tasks without blocking
-                    while True:
-                        try:
-                            task = tasks.receive_nowait()
-                            task_result = handle_task(task)
-                            if not task_result:
-                                running = False
-                                break
-                        except WouldBlock:
-                            break
-                    if not running:
-                        continue
-
-                # Sync and insert pending requests (collective operation)
-                # Rank 0 broadcasts its pending to all ranks
-                inserted = batch_engine.sync_and_insert_pending()
-                if is_rank_0 and inserted:
-                    current_status = RunnerRunning(
-                        active_requests=batch_engine.active_count
+                    logger.info(f"received chat request: {str(task)[:500]}")
+                    current_status = RunnerRunning()
+                    logger.info("runner running")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
                     )
-                    send_status(current_status)
+                    assert model
+                    assert tokenizer
+                    assert task_params.messages[0].content is not None
 
-                # Run generation for time budget
-                if batch_engine.has_active_requests:
-                    time_budget = TimeBudget(budget=0.5, group=group)
-                    for _ in time_budget:
-                        if not batch_engine.has_active_requests:
-                            break
-                        for resp in batch_engine.step():
-                            # Send token IMMEDIATELY for smooth streaming (only rank 0)
-                            if is_rank_0:
-                                event_sender.send(
-                                    ChunkGenerated(
-                                        command_id=resp.command_id,
-                                        chunk=TokenChunk(
-                                            idx=resp.response.token,
-                                            model=shard_metadata.model_meta.model_id,
-                                            text=resp.response.text,
-                                            token_id=resp.response.token,
-                                            finish_reason=resp.response.finish_reason,
-                                            stats=resp.response.stats,
-                                        ),
-                                    )
-                                )
-                                if resp.response.finish_reason is not None:
-                                    event_sender.send(
-                                        TaskStatusUpdated(
-                                            task_id=resp.task_id,
-                                            task_status=TaskStatus.Complete,
-                                        )
-                                    )
-
-                # Sync completions at budget boundary (always call - it's a collective operation)
-                batch_engine.sync_completions()
-
-                # Update status after budget
-                if is_rank_0:
-                    current_status = (
-                        RunnerRunning(active_requests=batch_engine.active_count)
-                        if batch_engine.has_active_requests
-                        else RunnerReady()
-                    )
-                    send_status(current_status)
-
-            else:
-                # Non-distributed mode: original logic with queue + insert
-                while True:
                     try:
-                        task = tasks.receive_nowait()
-                        running = handle_task(task)
-                        if not running:
-                            break
-                    except WouldBlock:
-                        break
+                        _check_for_debug_prompts(task_params.messages[0].content)
 
-                if not running:
-                    break
+                        # Build prompt once - used for both generation and thinking detection
+                        prompt = apply_chat_template(tokenizer, task_params)
 
-                # Insert any queued requests (non-distributed just inserts directly)
-                # Status was already sent in handle_task when queueing
-                if batch_engine is not None and batch_engine.has_pending_inserts:
-                    batch_engine.sync_and_insert_pending()
+                        # Generate responses using the actual MLX generation
+                        mlx_generator = mlx_generate(
+                            model=model,
+                            tokenizer=tokenizer,
+                            task=task_params,
+                            prompt=prompt,
+                        )
 
-                if batch_engine is not None and batch_engine.has_active_requests:
-                    for resp in batch_engine.step():
-                        if shard_metadata.device_rank == 0:
+                        # 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):
+                            mlx_generator = parse_thinking_models(
+                                mlx_generator, tokenizer
+                            )
+
+                        # TODO: Add tool call parser here
+
+                        for response in mlx_generator:
+                            match response:
+                                case GenerationResponse():
+                                    if device_rank == 0:
+                                        event_sender.send(
+                                            ChunkGenerated(
+                                                command_id=command_id,
+                                                chunk=TokenChunk(
+                                                    idx=response.token,
+                                                    model=shard_metadata.model_card.model_id,
+                                                    text=response.text,
+                                                    token_id=response.token,
+                                                    finish_reason=response.finish_reason,
+                                                    stats=response.stats,
+                                                ),
+                                            )
+                                        )
+
+                    # can we make this more explicit?
+                    except Exception as e:
+                        if device_rank == 0:
                             event_sender.send(
                                 ChunkGenerated(
-                                    command_id=resp.command_id,
+                                    command_id=command_id,
                                     chunk=TokenChunk(
-                                        idx=resp.response.token,
-                                        model=shard_metadata.model_meta.model_id,
-                                        text=resp.response.text,
-                                        token_id=resp.response.token,
-                                        finish_reason=resp.response.finish_reason,
-                                        stats=resp.response.stats,
+                                        idx=0,
+                                        model=shard_metadata.model_card.model_id,
+                                        text="",
+                                        token_id=0,
+                                        finish_reason="error",
+                                        error_message=str(e),
                                     ),
                                 )
                             )
-                        if resp.response.finish_reason is not None:
-                            event_sender.send(
-                                TaskStatusUpdated(
-                                    task_id=resp.task_id,
-                                    task_status=TaskStatus.Complete,
-                                )
-                            )
+                        raise
 
-                    if batch_engine.has_active_requests:
-                        current_status = RunnerRunning(
-                            active_requests=batch_engine.active_count
+                    current_status = RunnerReady()
+                    logger.info("runner ready")
+                case Shutdown():
+                    current_status = RunnerShuttingDown()
+                    logger.info("runner shutting down")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
                         )
-                    else:
-                        current_status = RunnerReady()
-                    send_status(current_status)
+                    )
+                    current_status = RunnerShutdown()
+                case _:
+                    raise ValueError(
+                        f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
+                    )
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
+            )
+            event_sender.send(
+                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+            )
+            if isinstance(current_status, RunnerShutdown):
+                del model, tokenizer, group
+                mx.clear_cache()
+                import gc
 
-                    # Process deferred shutdown after all requests complete
-                    if (
-                        pending_shutdown is not None
-                        and not batch_engine.has_active_requests
-                        and not batch_engine.has_pending_inserts
-                    ):
-                        running = handle_task(pending_shutdown, is_deferred=True)
-                else:
-                    task = tasks.receive()
-                    running = handle_task(task)
+                gc.collect()
+                break
 
-    # Cleanup
-    del model, tokenizer, group, batch_engine
-    mx.clear_cache()
-    gc.collect()
+
+@cache
+def get_gpt_oss_encoding():
+    encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+    return encoding
+
+
+def parse_gpt_oss(
+    responses: Generator[GenerationResponse],
+) -> Generator[GenerationResponse]:
+    encoding = get_gpt_oss_encoding()
+    stream = StreamableParser(encoding, role=Role.ASSISTANT)
+    thinking = False
+
+    for response in responses:
+        stream.process(response.token)
+
+        delta = stream.last_content_delta
+        ch = stream.current_channel
+
+        if ch == "analysis" and not thinking:
+            thinking = True
+            yield response.model_copy(update={"text": "<think>"})
+
+        if ch != "analysis" and thinking:
+            thinking = False
+            yield response.model_copy(update={"text": "</think>"})
+
+        if delta:
+            yield response.model_copy(update={"text": delta})
+
+        if response.finish_reason is not None:
+            if thinking:
+                yield response.model_copy(update={"text": "</think>"})
+            yield response
+            break
+
+
+def parse_thinking_models(
+    responses: Generator[GenerationResponse],
+    tokenizer: TokenizerWrapper,
+) -> 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
+    can properly parse thinking content.
+    """
+    first = True
+    for response in responses:
+        if first:
+            first = False
+            yield response.model_copy(
+                update={
+                    "text": tokenizer.think_start,
+                    "token": tokenizer.think_start_id,  # type: ignore
+                }
+            )
+        yield response
 
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"

src/exo/worker/runner/runner_supervisor.py

@@ -105,7 +105,7 @@ class RunnerSupervisor:
             return
 
         # This is overkill but it's not technically bad, just unnecessary.
-        logger.warning("Runner process didn't shutdown successfully, terminating")
+        logger.warning("Runner process didn't shutdown succesfully, terminating")
         self.runner_process.terminate()
         await to_thread.run_sync(self.runner_process.join, 5)
         if not self.runner_process.is_alive():
@@ -128,11 +128,9 @@ class RunnerSupervisor:
 
     async def start_task(self, task: Task):
         if task.task_id in self.completed:
-            logger.info(f"Skipping task {task.task_id} - already completed")
-            return
-        if task.task_id in self.pending:
-            logger.info(f"Skipping task {task.task_id} - already pending")
-            return
+            logger.info(
+                f"Skipping invalid task {task} as it has already been completed"
+            )
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
@@ -151,17 +149,13 @@ class RunnerSupervisor:
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
-                        # Just set the event to unblock start_task, but keep in pending
-                        # to prevent duplicate forwarding until completion
-                        if event.task_id in self.pending:
-                            self.pending[event.task_id].set()
+                        self.pending.pop(event.task_id).set()
                         continue
-                    if isinstance(event, TaskStatusUpdated) and event.task_status in (
-                        TaskStatus.Complete,
-                        TaskStatus.TimedOut,
-                        TaskStatus.Failed,
+                    if (
+                        isinstance(event, TaskStatusUpdated)
+                        and event.task_status == TaskStatus.Complete
                     ):
-                        # If a task has just finished, we should be working on it.
+                        # If a task has just been completed, we should be working on it.
                         assert isinstance(
                             self.status,
                             (
@@ -172,8 +166,6 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
-                        # Now safe to remove from pending and add to completed
-                        self.pending.pop(event.task_id, None)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
             except (ClosedResourceError, BrokenResourceError) as e:

src/exo/worker/tests/constants.py

@@ -1,7 +1,7 @@
 from typing import Final
 
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import TaskId
 from exo.shared.types.worker.instances import InstanceId, RunnerId
 

src/exo/worker/tests/unittests/conftest.py

@@ -1,8 +1,8 @@
 from dataclasses import dataclass, field
 
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.tasks import BaseTask, TaskId
 from exo.shared.types.worker.instances import (
     BoundInstance,
@@ -20,7 +20,6 @@ class FakeRunnerSupervisor:
     bound_instance: BoundInstance
     status: RunnerStatus
     completed: set[TaskId] = field(default_factory=set)
-    pending: dict[TaskId, object] = field(default_factory=dict)
 
 
 class OtherTask(BaseTask):
@@ -33,9 +32,8 @@ def get_pipeline_shard_metadata(
     model_id: ModelId, device_rank: int, world_size: int = 1
 ) -> ShardMetadata:
     return PipelineShardMetadata(
-        model_meta=ModelMetadata(
+        model_card=ModelCard(
             model_id=model_id,
-            pretty_name=str(model_id),
             storage_size=Memory.from_mb(100000),
             n_layers=32,
             hidden_size=2048,

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

@@ -0,0 +1,199 @@
+# type: ignore
+import json
+import os
+import tempfile
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.types.api import ChatCompletionMessage
+from exo.shared.types.memory import Memory
+from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.generator.generate import mlx_generate
+from exo.worker.engines.mlx.utils_mlx import shard_and_load
+
+
+class MockLayer(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.custom_attr = "test_value"
+        self.use_sliding = True
+
+    def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+        return x * 2
+
+
+@dataclass(frozen=True)
+class PipelineTestConfig:
+    model_path: Path
+    total_layers: int
+    base_port: int
+    max_tokens: int
+
+
+def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]:
+    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        hostfile_path = f.name
+
+    return hostfile_path, hosts
+
+
+# Use GPT OSS 20b to test as it is a model with a lot of strange behaviour
+
+DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig(
+    model_path=EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8",
+    total_layers=24,
+    base_port=29600,
+    max_tokens=200,
+)
+
+
+DEFAULT_GPT_OSS_MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+
+
+def run_gpt_oss_pipeline_device(
+    rank: int,
+    world_size: int,
+    hostfile_path: str,
+    layer_splits: list[tuple[int, int]],
+    prompt_tokens: int,
+    prefill_step_size: int,
+    result_queue: Any,  # pyright: ignore[reportAny]
+    max_tokens: int = 200,
+) -> None:
+    os.environ["MLX_HOSTFILE"] = hostfile_path
+    os.environ["MLX_RANK"] = str(rank)
+
+    try:
+        group = mx.distributed.init(backend="ring", strict=True)
+
+        start_layer, end_layer = layer_splits[rank]
+
+        shard_meta = PipelineShardMetadata(
+            model_card=ModelCard(
+                model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
+                storage_size=Memory.from_gb(12),
+                n_layers=24,
+                hidden_size=2880,
+                supports_tensor=False,
+            ),
+            device_rank=rank,
+            world_size=world_size,
+            start_layer=start_layer,
+            end_layer=end_layer,
+            n_layers=24,
+        )
+
+        model, tokenizer = shard_and_load(shard_meta, group)
+        model = cast(Model, model)
+
+        # Generate a prompt of exact token length
+        base_text = "The quick brown fox jumps over the lazy dog. "
+        base_tokens = tokenizer.encode(base_text)
+        base_len = len(base_tokens)
+
+        # Build prompt with approximate target length
+        repeats = (prompt_tokens // base_len) + 2
+        long_text = base_text * repeats
+        tokens = tokenizer.encode(long_text)
+        # Truncate to exact target length
+        tokens = tokens[:prompt_tokens]
+        prompt_text = tokenizer.decode(tokens)
+
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content=prompt_text)],
+            max_tokens=max_tokens,
+        )
+
+        generated_text = ""
+        for response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+        ):
+            generated_text += response.text
+            if response.finish_reason is not None:
+                break
+
+        result_queue.put((rank, True, generated_text))  # pyright: ignore[reportAny]
+
+    except Exception as e:
+        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))  # pyright: ignore[reportAny]
+
+
+def run_gpt_oss_tensor_parallel_device(
+    rank: int,
+    world_size: int,
+    hostfile_path: str,
+    prompt_tokens: int,
+    prefill_step_size: int,
+    result_queue: Any,  # pyright: ignore[reportAny]
+    max_tokens: int = 10,
+) -> None:
+    os.environ["MLX_HOSTFILE"] = hostfile_path
+    os.environ["MLX_RANK"] = str(rank)
+
+    try:
+        group = mx.distributed.init(backend="ring", strict=True)
+
+        # For tensor parallelism, all devices run all layers
+        shard_meta = TensorShardMetadata(
+            model_card=ModelCard(
+                model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
+                storage_size=Memory.from_gb(12),
+                n_layers=24,
+                hidden_size=2880,
+                supports_tensor=True,
+            ),
+            device_rank=rank,
+            world_size=world_size,
+            start_layer=0,
+            end_layer=24,
+            n_layers=24,
+        )
+
+        model, tokenizer = shard_and_load(shard_meta, group)
+        model = cast(Model, model)
+
+        base_text = "The quick brown fox jumps over the lazy dog. "
+        base_tokens = tokenizer.encode(base_text)
+        base_len = len(base_tokens)
+
+        repeats = (prompt_tokens // base_len) + 2
+        long_text = base_text * repeats
+        tokens = tokenizer.encode(long_text)
+        tokens = tokens[:prompt_tokens]
+        prompt_text = tokenizer.decode(tokens)
+
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content=prompt_text)],
+            max_tokens=max_tokens,
+        )
+
+        generated_text = ""
+        for response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+        ):
+            generated_text += response.text
+            if response.finish_reason is not None:
+                break
+
+        result_queue.put((rank, True, generated_text))  # pyright: ignore[reportAny]
+
+    except Exception as e:
+        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))  # pyright: ignore[reportAny]

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

@@ -0,0 +1,146 @@
+import json
+import multiprocessing as mp
+import os
+import tempfile
+from typing import Any
+
+import mlx.core as mx
+import mlx.nn as mlx_nn
+import pytest
+
+from exo.worker.engines.mlx.auto_parallel import (
+    CustomMlxLayer,
+    PipelineFirstLayer,
+    PipelineLastLayer,
+    patch_pipeline_model,
+)
+from exo.worker.tests.unittests.test_mlx.conftest import MockLayer
+
+
+def run_pipeline_device(
+    rank: int,
+    world_size: int,
+    hostfile_path: str,
+    result_queue: Any,  # pyright: ignore[reportAny]
+) -> None:
+    import os
+
+    os.environ["MLX_HOSTFILE"] = hostfile_path
+    os.environ["MLX_RANK"] = str(rank)
+
+    class MockLayerInner(mlx_nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.custom_attr = "test_value"
+
+        def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+            return x * 2
+
+    class MockModel(mlx_nn.Module):
+        def __init__(self, layers: list[mlx_nn.Module]) -> None:
+            super().__init__()
+            self.layers = layers
+
+        def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+            for layer in self.layers:
+                x = layer(x, *args, **kwargs)  # pyright: ignore[reportUnknownVariableType]
+            return x  # pyright: ignore[reportUnknownVariableType]
+
+    try:
+        group = mx.distributed.init(backend="ring", strict=True)
+
+        mock = MockLayerInner()
+        first = PipelineFirstLayer(mock, r=rank, group=group)
+        composed = PipelineLastLayer(first, r=rank, s=world_size, group=group)
+
+        # Wrap in a mock model, then wrap in PipelineParallelModel for all_gather
+        inner_model = MockModel([composed])
+        model = patch_pipeline_model(inner_model, group)
+
+        x = mx.ones((1, 4))
+        result = model(x)
+        mx.eval(result)
+        success = result.shape == x.shape
+        result_queue.put((rank, success, result))  # pyright: ignore[reportAny]
+    except Exception as e:
+        result_queue.put((rank, False, str(e)))  # pyright: ignore[reportAny]
+
+
+def test_single_wrapper_delegates_attributes() -> None:
+    mock = MockLayer()
+    wrapped = CustomMlxLayer(mock)
+
+    assert wrapped.custom_attr == "test_value"  # type: ignore[attr-defined]
+    assert wrapped.use_sliding is True  # type: ignore[attr-defined]
+
+
+def test_composed_wrappers_delegate_attributes() -> None:
+    mock = MockLayer()
+    group = mx.distributed.init()
+
+    first = PipelineFirstLayer(mock, r=0, group=group)
+    composed = PipelineLastLayer(first, r=0, s=1, group=group)
+
+    assert composed.custom_attr == "test_value"  # type: ignore[attr-defined]
+    assert composed.use_sliding is True  # type: ignore[attr-defined]
+
+
+def test_missing_attribute_raises() -> None:
+    mock = MockLayer()
+    wrapped = CustomMlxLayer(mock)
+
+    with pytest.raises(AttributeError):
+        _ = wrapped.nonexistent_attr  # type: ignore[attr-defined]
+
+
+def test_composed_call_works() -> None:
+    ctx = mp.get_context("spawn")
+
+    world_size = 2
+    base_port = 29500
+
+    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        hostfile_path = f.name
+
+    try:
+        result_queue: Any = ctx.Queue()
+
+        processes: list[Any] = []
+        for rank in range(world_size):
+            p = ctx.Process(
+                target=run_pipeline_device,
+                args=(rank, world_size, hostfile_path, result_queue),
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:  # pyright: ignore[reportAny]
+            p.join(timeout=10)  # pyright: ignore[reportAny]
+
+        results: dict[int, Any] = {}
+        errors: dict[int, str] = {}
+        while not result_queue.empty():  # pyright: ignore[reportAny]
+            rank, success, value = result_queue.get()  # pyright: ignore[reportAny]
+            if success:
+                results[rank] = value
+            else:
+                errors[rank] = value
+
+        assert len(results) == world_size, (
+            f"Expected {world_size} results, got {len(results)}. Errors: {errors}"
+        )
+
+        for rank in range(world_size):
+            assert rank in results, (
+                f"Device {rank} failed: {errors.get(rank, 'unknown')}"
+            )
+            result_array = results[rank]
+            # Both devices see the final result (4.0) after all_gather
+            assert (result_array == 4.0).all(), (
+                f"Device {rank}: expected 4.0, got {result_array}"
+            )
+    finally:
+        os.unlink(hostfile_path)

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

@@ -0,0 +1,231 @@
+import multiprocessing as mp
+import os
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import pytest
+
+from exo.worker.tests.unittests.test_mlx.conftest import (
+    DEFAULT_GPT_OSS_CONFIG,
+    create_hostfile,
+    run_gpt_oss_pipeline_device,
+    run_gpt_oss_tensor_parallel_device,
+)
+
+
+def _check_model_exists() -> bool:
+    return DEFAULT_GPT_OSS_CONFIG.model_path.exists()
+
+
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(
+        not _check_model_exists(),
+        reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}",
+    ),
+]
+
+
+@dataclass
+class DistributedTestResult:
+    timed_out: bool
+    world_size: int
+    results: dict[int, tuple[bool, str]]
+
+    @property
+    def all_success(self) -> bool:
+        if len(self.results) != self.world_size:
+            return False
+        return all(r[0] for r in self.results.values())
+
+
+def run_distributed_test(
+    world_size: int,
+    port_offset: int,
+    process_timeout: int,
+    target: Callable[..., None],
+    make_args: Callable[[int], tuple[Any, ...]],
+) -> DistributedTestResult:
+    ctx = mp.get_context("spawn")
+    hostfile_path, _ = create_hostfile(
+        world_size, DEFAULT_GPT_OSS_CONFIG.base_port + port_offset
+    )
+
+    try:
+        result_queue: Any = ctx.Queue()
+        processes: list[Any] = []
+
+        for rank in range(world_size):
+            args = make_args(rank)
+            p = ctx.Process(
+                target=target,
+                args=(rank, world_size, hostfile_path, *args, result_queue),
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:  # pyright: ignore[reportAny]
+            p.join(timeout=process_timeout)  # pyright: ignore[reportAny]
+
+        timed_out = any(p.is_alive() for p in processes)  # pyright: ignore[reportAny]
+
+        for p in processes:  # pyright: ignore[reportAny]
+            if p.is_alive():  # pyright: ignore[reportAny]
+                p.terminate()  # pyright: ignore[reportAny]
+                p.join(timeout=5)  # pyright: ignore[reportAny]
+
+        results: dict[int, tuple[bool, str]] = {}
+        while not result_queue.empty():  # pyright: ignore[reportAny]
+            rank, success, value = result_queue.get()  # pyright: ignore[reportAny]
+            results[rank] = (success, value)
+
+        return DistributedTestResult(
+            timed_out=timed_out, world_size=world_size, results=results
+        )
+
+    finally:
+        os.unlink(hostfile_path)
+
+
+def run_pipeline_test(
+    layer_splits: list[tuple[int, int]],
+    prompt_tokens: int,
+    prefill_step_size: int,
+    port_offset: int = 0,
+    process_timeout: int = 60,
+) -> DistributedTestResult:
+    def make_args(rank: int) -> tuple[Any, ...]:
+        return (
+            layer_splits,
+            prompt_tokens,
+            prefill_step_size,
+        )
+
+    return run_distributed_test(
+        world_size=len(layer_splits),
+        port_offset=port_offset,
+        process_timeout=process_timeout,
+        target=run_gpt_oss_pipeline_device,
+        make_args=make_args,
+    )
+
+
+def run_tensor_test(
+    prompt_tokens: int,
+    prefill_step_size: int,
+    port_offset: int = 0,
+    process_timeout: int = 60,
+) -> DistributedTestResult:
+    def make_args(rank: int) -> tuple[Any, ...]:
+        return (
+            prompt_tokens,
+            prefill_step_size,
+        )
+
+    return run_distributed_test(
+        world_size=2,
+        port_offset=port_offset,
+        process_timeout=process_timeout,
+        target=run_gpt_oss_tensor_parallel_device,
+        make_args=make_args,
+    )
+
+
+class TestPipelineParallelFix:
+    BUG_TRIGGER_SPLITS: list[tuple[int, int]] = [(0, 1), (1, 24)]
+
+    def test_pipeline_single_layer_first_device(self) -> None:
+        result = run_pipeline_test(
+            layer_splits=self.BUG_TRIGGER_SPLITS,
+            prompt_tokens=100,
+            prefill_step_size=64,
+            process_timeout=60,
+        )
+        assert not result.timed_out, "Unexpected timeout - fix may not be working"
+        assert result.all_success, f"Failures: {result.results}"
+
+
+class TestPipelineSplitConfigurations:
+    @pytest.mark.parametrize(
+        "layer_splits",
+        [
+            [(0, 1), (1, 24)],
+            [(0, 6), (6, 24)],
+            [(0, 12), (12, 24)],
+        ],
+        ids=["1_23", "6_18", "12_12"],
+    )
+    def test_pipeline_splits(
+        self,
+        layer_splits: list[tuple[int, int]],
+    ) -> None:
+        result = run_pipeline_test(
+            layer_splits=layer_splits,
+            prompt_tokens=600,
+            prefill_step_size=512,
+            port_offset=100,
+        )
+        assert not result.timed_out, f"Timeout with {layer_splits}"
+        assert result.all_success, f"Failures with {layer_splits}: {result.results}"
+
+
+class TestPrefillStepSizeBoundaries:
+    @pytest.mark.parametrize(
+        "prefill_step_size,prompt_tokens",
+        [
+            (512, 511),
+            (512, 512),
+            (512, 513),
+            (512, 1024),
+        ],
+        ids=["under", "exact", "over", "double"],
+    )
+    def test_boundary_conditions(
+        self,
+        prefill_step_size: int,
+        prompt_tokens: int,
+    ) -> None:
+        result = run_pipeline_test(
+            layer_splits=[(0, 12), (12, 24)],
+            prompt_tokens=prompt_tokens,
+            prefill_step_size=prefill_step_size,
+            port_offset=200,
+        )
+        assert not result.timed_out, f"Timeout: {prompt_tokens=}, {prefill_step_size=}"
+        assert result.all_success, f"Failures: {result.results}"
+
+
+class TestTensorParallelFix:
+    def test_tensor_parallel(self) -> None:
+        result = run_tensor_test(
+            prompt_tokens=100,
+            prefill_step_size=64,
+            port_offset=400,
+        )
+        assert not result.timed_out, "Unexpected timeout"
+        assert result.all_success, f"Failures: {result.results}"
+
+
+class TestTensorParallelBoundaries:
+    @pytest.mark.parametrize(
+        "prefill_step_size,prompt_tokens",
+        [
+            (512, 511),
+            (512, 512),
+            (512, 513),
+            (512, 1024),
+        ],
+        ids=["under", "exact", "over", "double"],
+    )
+    def test_tensor_parallel_boundaries(
+        self,
+        prefill_step_size: int,
+        prompt_tokens: int,
+    ) -> None:
+        result = run_tensor_test(
+            prompt_tokens=prompt_tokens,
+            prefill_step_size=prefill_step_size,
+            port_offset=500,
+        )
+        assert not result.timed_out, f"Timeout: {prompt_tokens=}, {prefill_step_size=}"
+        assert result.all_success, f"Failures: {result.results}"

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

@@ -76,19 +76,21 @@ def get_test_models() -> list[tuple[str, ModelCard]]:
     """Get a representative sample of models to test."""
     # Pick one model from each family to test
     families: dict[str, tuple[str, ModelCard]] = {}
-    for short_id, card in MODEL_CARDS.items():
+    for _, card in MODEL_CARDS.items():
         # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
-        parts = short_id.split("-")
+        parts = card.model_id.short().split("-")
         family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
 
         if family not in families:
-            families[family] = (short_id, card)
+            families[family] = (card.model_id.short(), card)
 
     return list(families.values())
 
 
 TEST_MODELS: list[tuple[str, ModelCard]] = get_test_models()
 
+pytestmark = pytest.mark.slow
+
 
 @pytest.fixture(scope="module")
 def event_loop():

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

@@ -1,7 +1,7 @@
 import exo.worker.plan as plan_mod
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import LoadModel
 from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
 from exo.shared.types.worker.instances import BoundInstance

src/exo/worker/tests/unittests/test_runner/test_continuous_batching.py

@@ -1,319 +0,0 @@
-"""
-Tests for continuous batching behavior in the runner.
-
-These tests verify that:
-1. Single requests work through the batch path
-2. Multiple concurrent requests batch together
-3. Tokens are routed to the correct requests
-4. Requests complete at different times appropriately
-"""
-
-# pyright: reportAny=false
-# pyright: reportUnknownArgumentType=false
-# pyright: reportUnknownMemberType=false
-# pyright: reportAttributeAccessIssue=false
-# pyright: reportInvalidTypeVarUse=false
-
-from typing import Any
-from unittest.mock import MagicMock
-
-import pytest
-
-import exo.worker.runner.runner as mlx_runner
-from exo.shared.types.api import ChatCompletionMessage
-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.events import (
-    Event,
-    RunnerStatusUpdated,
-    TaskStatusUpdated,
-)
-from exo.shared.types.tasks import (
-    ChatCompletion,
-    ChatCompletionTaskParams,
-    ConnectToGroup,
-    LoadModel,
-    Shutdown,
-    StartWarmup,
-    Task,
-    TaskId,
-    TaskStatus,
-)
-from exo.shared.types.worker.runner_response import GenerationResponse
-from exo.shared.types.worker.runners import RunnerRunning
-from exo.utils.channels import mp_channel
-from exo.worker.engines.mlx.generator.batch_engine import (
-    BatchedGenerationResponse,
-)
-from exo.worker.tests.constants import (
-    INSTANCE_1_ID,
-    MODEL_A_ID,
-    NODE_A,
-    RUNNER_1_ID,
-)
-from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
-
-
-class FakeBatchEngineWithTokens:
-    """
-    Fake batch engine that generates a specified number of tokens per request.
-
-    This simulates realistic batch generation behavior where:
-    - Requests are queued on insert
-    - Each step() call generates one token for all active requests
-    - Requests complete when they've generated all their tokens
-    """
-
-    def __init__(self, *_args: Any, **_kwargs: Any):
-        self._active_requests: dict[int, tuple[CommandId, TaskId, int, int]] = {}
-        self._pending_inserts: list[
-            tuple[CommandId, TaskId, ChatCompletionTaskParams]
-        ] = []
-        self._uid_counter = 0
-        self._tokens_per_request = 3  # Default: generate 3 tokens before completing
-        self.rank = 0  # Fake rank for testing
-
-    def queue_request(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams,
-    ) -> None:
-        """Queue a request for insertion."""
-        self._pending_inserts.append((command_id, task_id, task_params))
-
-    def sync_and_insert_pending(self) -> list[int]:
-        """Insert all pending requests."""
-        uids: list[int] = []
-        for command_id, task_id, task_params in self._pending_inserts:
-            uid = self._do_insert(command_id, task_id, task_params)
-            uids.append(uid)
-        self._pending_inserts.clear()
-        return uids
-
-    @property
-    def has_pending_inserts(self) -> bool:
-        return len(self._pending_inserts) > 0
-
-    def _do_insert(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams | None,
-    ) -> int:
-        uid = self._uid_counter
-        self._uid_counter += 1
-        # Track: (command_id, task_id, tokens_generated, max_tokens)
-        max_tokens = task_params.max_tokens if task_params else self._tokens_per_request
-        self._active_requests[uid] = (command_id, task_id, 0, max_tokens or 3)
-        return uid
-
-    def step(self) -> list[BatchedGenerationResponse]:
-        results: list[BatchedGenerationResponse] = []
-        uids_to_remove: list[int] = []
-
-        for uid, (command_id, task_id, tokens_gen, max_tokens) in list(
-            self._active_requests.items()
-        ):
-            tokens_gen += 1
-            finish_reason = "stop" if tokens_gen >= max_tokens else None
-            text = f"token{tokens_gen}"
-
-            if finish_reason:
-                uids_to_remove.append(uid)
-            else:
-                self._active_requests[uid] = (
-                    command_id,
-                    task_id,
-                    tokens_gen,
-                    max_tokens,
-                )
-
-            results.append(
-                BatchedGenerationResponse(
-                    command_id=command_id,
-                    task_id=task_id,
-                    response=GenerationResponse(
-                        token=tokens_gen,
-                        text=text,
-                        finish_reason=finish_reason,
-                    ),
-                )
-            )
-
-        for uid in uids_to_remove:
-            del self._active_requests[uid]
-
-        return results
-
-    @property
-    def has_active_requests(self) -> bool:
-        return len(self._active_requests) > 0
-
-    @property
-    def active_count(self) -> int:
-        return len(self._active_requests)
-
-    @property
-    def pending_insert_count(self) -> int:
-        return len(self._pending_inserts)
-
-
-def make_nothin[T, U, V](res: T):
-    def nothin(*_1: U, **_2: V) -> T:
-        return res
-
-    return nothin
-
-
-@pytest.fixture
-def patch_batch_engine(monkeypatch: pytest.MonkeyPatch):
-    """Patch MLX dependencies and use FakeBatchEngineWithTokens."""
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MagicMock()))
-    monkeypatch.setattr(
-        mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MagicMock()))
-    )
-    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
-    monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", make_nothin(None))
-    monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngineWithTokens)
-
-
-def _run_with_tasks(tasks: list[Task]) -> list[Event]:
-    """
-    Run tasks through the runner, adding shutdown at the end.
-
-    Tasks are sent in order, with shutdown sent last.
-    The batch engine processes between task handling.
-    """
-    bound_instance = get_bound_mlx_ring_instance(
-        instance_id=INSTANCE_1_ID,
-        model_id=MODEL_A_ID,
-        runner_id=RUNNER_1_ID,
-        node_id=NodeId(NODE_A),
-    )
-
-    task_sender, task_receiver = mp_channel[Task]()
-    event_sender, event_receiver = mp_channel[Event]()
-
-    shutdown_task = Shutdown(
-        task_id=TaskId("shutdown"),
-        instance_id=INSTANCE_1_ID,
-        runner_id=RUNNER_1_ID,
-    )
-
-    with task_sender, event_receiver:
-        # Send all tasks including shutdown
-        for t in tasks:
-            task_sender.send(t)
-        task_sender.send(shutdown_task)
-
-        # Disable cleanup methods to prevent issues
-        event_sender.close = lambda: None
-        event_sender.join = lambda: None
-        task_receiver.close = lambda: None
-        task_receiver.join = lambda: None
-
-        mlx_runner.main(bound_instance, event_sender, task_receiver)
-
-        return event_receiver.collect()
-
-
-INIT_TASK = ConnectToGroup(task_id=TaskId("init"), instance_id=INSTANCE_1_ID)
-LOAD_TASK = LoadModel(task_id=TaskId("load"), instance_id=INSTANCE_1_ID)
-WARMUP_TASK = StartWarmup(task_id=TaskId("warmup"), instance_id=INSTANCE_1_ID)
-
-
-def make_chat_task(
-    task_id: str, command_id: str, max_tokens: int = 3
-) -> ChatCompletion:
-    return ChatCompletion(
-        task_id=TaskId(task_id),
-        command_id=CommandId(command_id),
-        task_params=ChatCompletionTaskParams(
-            model=str(MODEL_A_ID),
-            messages=[ChatCompletionMessage(role="user", content="hello")],
-            stream=True,
-            max_tokens=max_tokens,
-        ),
-        instance_id=INSTANCE_1_ID,
-    )
-
-
-def test_single_request_generates_tokens(patch_batch_engine: None):
-    """
-    Verify a single request generates the expected tokens through the batch path.
-
-    Note: With the current non-blocking design, shutdown is processed before
-    batch steps run when all tasks are queued together. This test verifies
-    the runner status reflects active requests.
-    """
-    chat_task = make_chat_task("chat1", "cmd1", max_tokens=3)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
-
-    # Find RunnerRunning status events - this shows the request was inserted
-    running_events = [
-        e
-        for e in events
-        if isinstance(e, RunnerStatusUpdated)
-        and isinstance(e.runner_status, RunnerRunning)
-    ]
-
-    assert len(running_events) >= 1, "Expected at least one RunnerRunning event"
-    assert running_events[0].runner_status.active_requests == 1
-
-
-def test_runner_status_reflects_active_requests(patch_batch_engine: None):
-    """Verify RunnerRunning status includes active_requests count."""
-    chat_task = make_chat_task("chat1", "cmd1", max_tokens=2)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
-
-    # Find RunnerRunning status events
-    running_events = [
-        e
-        for e in events
-        if isinstance(e, RunnerStatusUpdated)
-        and isinstance(e.runner_status, RunnerRunning)
-    ]
-
-    assert len(running_events) > 0, "Expected at least one RunnerRunning event"
-    assert running_events[0].runner_status.active_requests == 1
-
-
-def test_chat_task_acknowledged(patch_batch_engine: None):
-    """Verify chat completion task is acknowledged with proper status updates."""
-    chat_task = make_chat_task("chat1", "cmd1", max_tokens=2)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
-
-    # Find the chat task status events
-    chat_running = [
-        e
-        for e in events
-        if isinstance(e, TaskStatusUpdated)
-        and e.task_id == TaskId("chat1")
-        and e.task_status == TaskStatus.Running
-    ]
-
-    assert len(chat_running) == 1, "Expected exactly one chat task Running status"
-
-
-def test_multiple_requests_tracked(patch_batch_engine: None):
-    """Verify multiple concurrent requests are tracked in active_requests."""
-    chat1 = make_chat_task("chat1", "cmd1", max_tokens=2)
-    chat2 = make_chat_task("chat2", "cmd2", max_tokens=2)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat1, chat2])
-
-    # Find RunnerRunning status events
-    running_events = [
-        e
-        for e in events
-        if isinstance(e, RunnerStatusUpdated)
-        and isinstance(e.runner_status, RunnerRunning)
-    ]
-
-    # Should have at least 2 RunnerRunning events (one per request inserted)
-    assert len(running_events) >= 2, (
-        f"Expected at least 2 RunnerRunning events, got {len(running_events)}"
-    )
-
-    # First should have 1 active request, second should have 2
-    assert running_events[0].runner_status.active_requests == 1
-    assert running_events[1].runner_status.active_requests == 2

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -1,17 +1,12 @@
 # Check tasks are complete before runner is ever ready.
-
-# pyright: reportAny=false
-
 from collections.abc import Iterable
-from typing import Any, Callable
-from unittest.mock import MagicMock
+from typing import Callable
 
 import pytest
 
 import exo.worker.runner.runner as mlx_runner
 from exo.shared.types.api import ChatCompletionMessage
 from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.common import CommandId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -27,7 +22,6 @@ from exo.shared.types.tasks import (
     Shutdown,
     StartWarmup,
     Task,
-    TaskId,
     TaskStatus,
 )
 from exo.shared.types.worker.runner_response import GenerationResponse
@@ -44,9 +38,6 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import mp_channel
-from exo.worker.engines.mlx.generator.batch_engine import (
-    BatchedGenerationResponse,
-)
 
 from ...constants import (
     CHAT_COMPLETION_TASK_ID,
@@ -116,89 +107,37 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even
         assert test_event == true_event, f"{test_event} != {true_event}"
 
 
-class FakeBatchEngine:
-    """
-    Fake batch engine for testing.
-
-    Queues requests on insert, returns one token per step.
-    The runner's non-blocking loop drains all tasks before running batch steps,
-    so this engine queues requests and has_active_requests returns True only
-    after at least one request has been inserted.
-    """
-
-    def __init__(self, *_args: Any, **_kwargs: Any):
-        self._active_requests: dict[int, tuple[CommandId, TaskId]] = {}
-        self._pending_inserts: list[
-            tuple[CommandId, TaskId, ChatCompletionTaskParams]
-        ] = []
-        self._uid_counter = 0
-        self.rank = 0  # Fake rank for testing
-
-    def queue_request(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams,
-    ) -> None:
-        """Queue a request for insertion."""
-        self._pending_inserts.append((command_id, task_id, task_params))
-
-    def sync_and_insert_pending(self) -> list[int]:
-        """Insert all pending requests."""
-        uids: list[int] = []
-        for command_id, task_id, _task_params in self._pending_inserts:
-            uid = self._uid_counter
-            self._uid_counter += 1
-            self._active_requests[uid] = (command_id, task_id)
-            uids.append(uid)
-        self._pending_inserts.clear()
-        return uids
-
-    @property
-    def has_pending_inserts(self) -> bool:
-        return len(self._pending_inserts) > 0
-
-    def step(self) -> list[BatchedGenerationResponse]:
-        results: list[BatchedGenerationResponse] = []
-        # Process all active requests - return one token and complete
-        for uid, (command_id, task_id) in list(self._active_requests.items()):
-            results.append(
-                BatchedGenerationResponse(
-                    command_id=command_id,
-                    task_id=task_id,
-                    response=GenerationResponse(
-                        token=0,
-                        text="hi",
-                        finish_reason="stop",
-                    ),
-                )
-            )
-            del self._active_requests[uid]
-        return results
-
-    @property
-    def has_active_requests(self) -> bool:
-        return len(self._active_requests) > 0
-
-    @property
-    def active_count(self) -> int:
-        return len(self._active_requests)
-
-    @property
-    def pending_insert_count(self) -> int:
-        return len(self._pending_inserts)
-
-
 @pytest.fixture
 def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
-    # initialize_mlx returns a fake "group" (non-None for state machine)
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MagicMock()))
-    monkeypatch.setattr(
-        mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MagicMock()))
-    )
+    # initialize_mlx returns a "group" equal to 1
+    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(1))
+    monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, 1)))
     monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
-    monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngine)
+    # Mock apply_chat_template since we're using a fake tokenizer (integer 1).
+    # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None.
+    monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt"))
+    monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
+
+    def fake_generate(*_1: object, **_2: object):
+        yield GenerationResponse(token=0, text="hi", finish_reason="stop")
+
+    monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
+
+
+# Use a fake event_sender to remove test flakiness.
+class EventCollector:
+    def __init__(self) -> None:
+        self.events: list[Event] = []
+
+    def send(self, event: Event) -> None:
+        self.events.append(event)
+
+    def close(self) -> None:
+        pass
+
+    def join(self) -> None:
+        pass
 
 
 def _run(tasks: Iterable[Task]):
@@ -210,26 +149,23 @@ def _run(tasks: Iterable[Task]):
     )
 
     task_sender, task_receiver = mp_channel[Task]()
-    event_sender, event_receiver = mp_channel[Event]()
+    event_sender = EventCollector()
 
-    with task_sender, event_receiver:
+    with task_sender:
         for t in tasks:
             task_sender.send(t)
 
         # worst monkeypatch known to man
         # this is some c++ nonsense
-        event_sender.close = nothin
-        event_sender.join = nothin
         task_receiver.close = nothin
         task_receiver.join = nothin
 
-        mlx_runner.main(bound_instance, event_sender, task_receiver)
+        mlx_runner.main(bound_instance, event_sender, task_receiver)  # type: ignore[arg-type]
 
-        return event_receiver.collect()
+        return event_sender.events
 
 
-def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPatch):
-    """Verify chat completion generates tokens, completes, and runner returns to Ready."""
+def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
     events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK, SHUTDOWN_TASK])
 
     expected_chunk = ChunkGenerated(
@@ -272,9 +208,7 @@ def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPat
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Running
             ),
             TaskAcknowledged(task_id=CHAT_COMPLETION_TASK_ID),
-            RunnerStatusUpdated(
-                runner_id=RUNNER_1_ID, runner_status=RunnerRunning(active_requests=1)
-            ),
+            RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerRunning()),
             expected_chunk,
             TaskStatusUpdated(
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Complete
@@ -289,6 +223,7 @@ def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPat
             TaskStatusUpdated(
                 task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete
             ),
+            # SPECIAL EXCEPTION FOR RUNNER SHUTDOWN
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()),
         ],
     )

src/exo/worker/utils/__init__.py

@@ -1,6 +0,0 @@
-from .profile import start_polling_memory_metrics, start_polling_node_metrics
-
-__all__ = [
-    "start_polling_node_metrics",
-    "start_polling_memory_metrics",
-]

src/exo/worker/utils/macmon.py

@@ -1,103 +0,0 @@
-import platform
-import shutil
-from subprocess import CalledProcessError
-from typing import cast
-
-from anyio import run_process
-from pydantic import BaseModel, ConfigDict, ValidationError
-
-
-class MacMonError(Exception):
-    """Exception raised for errors in the MacMon functions."""
-
-
-def _get_binary_path() -> str:
-    """
-    Get the path to the macmon binary.
-
-    Raises:
-        MacMonError: If the binary doesn't exist or can't be made executable.
-    """
-    # Check for macOS with ARM chip
-    system = platform.system().lower()
-    machine = platform.machine().lower()
-
-    if system != "darwin" or not (
-        "arm" in machine or "m1" in machine or "m2" in machine
-    ):
-        raise MacMonError("MacMon only supports macOS with Apple Silicon (ARM) chips")
-
-    path = shutil.which("macmon")
-
-    if path is None:
-        raise MacMonError("MacMon not found in PATH")
-
-    return path
-
-
-class TempMetrics(BaseModel):
-    """Temperature-related metrics returned by macmon."""
-
-    cpu_temp_avg: float
-    gpu_temp_avg: float
-
-    model_config = ConfigDict(extra="ignore")
-
-
-class Metrics(BaseModel):
-    """Complete set of metrics returned by macmon.
-
-    Unknown fields are ignored for forward-compatibility.
-    """
-
-    all_power: float
-    ane_power: float
-    cpu_power: float
-    ecpu_usage: tuple[int, float]
-    gpu_power: float
-    gpu_ram_power: float
-    gpu_usage: tuple[int, float]
-    pcpu_usage: tuple[int, float]
-    ram_power: float
-    sys_power: float
-    temp: TempMetrics
-    timestamp: str
-
-    model_config = ConfigDict(extra="ignore")
-
-
-async def get_metrics_async() -> Metrics:
-    """
-    Asynchronously run the binary and return the metrics as a Python dictionary.
-
-    Args:
-        binary_path: Optional path to the binary. If not provided, will use the bundled binary.
-
-    Returns:
-        A mapping containing system metrics.
-
-    Raises:
-        MacMonError: If there's an error running the binary.
-    """
-    path = _get_binary_path()
-
-    try:
-        # TODO: Keep Macmon running in the background?
-        result = await run_process([path, "pipe", "-s", "1"])
-
-        return Metrics.model_validate_json(result.stdout.decode().strip())
-
-    except ValidationError as e:
-        raise MacMonError(f"Error parsing JSON output: {e}") from e
-    except CalledProcessError as e:
-        stderr_msg = "no stderr"
-        stderr_output = cast(bytes | str | None, e.stderr)
-        if stderr_output is not None:
-            stderr_msg = (
-                stderr_output.decode()
-                if isinstance(stderr_output, bytes)
-                else str(stderr_output)
-            )
-        raise MacMonError(
-            f"MacMon failed with return code {e.returncode}: {stderr_msg}"
-        ) from e

src/exo/worker/utils/net_profile.py

@@ -1,91 +0,0 @@
-import http.client
-import time
-
-from anyio import create_task_group, to_thread
-from loguru import logger
-
-from exo.shared.topology import Topology
-from exo.shared.types.common import NodeId
-
-BAD_STATUSLINE_ATTEMPTS = 3
-
-
-async def check_reachability(
-    target_ip: str,
-    expected_node_id: NodeId,
-    self_node_id: NodeId,
-    out: dict[NodeId, set[str]],
-) -> None:
-    """Check if a node is reachable at the given IP and verify its identity."""
-
-    # TODO: use an async http client
-    def _fetch_remote_node_id(*, attempt: int = 1) -> NodeId | None:
-        connection = http.client.HTTPConnection(target_ip, 52415, timeout=3)
-        try:
-            connection.request("GET", "/node_id")
-            response = connection.getresponse()
-            if response.status != 200:
-                return None
-
-            body = response.read().decode("utf-8").strip()
-
-            # Strip quotes if present (JSON string response)
-            if body.startswith('"') and body.endswith('"') and len(body) >= 2:
-                body = body[1:-1]
-
-            return NodeId(body) or None
-        except OSError:
-            return None
-        except http.client.BadStatusLine:
-            if attempt >= BAD_STATUSLINE_ATTEMPTS:
-                logger.warning(
-                    f"BadStatusLine from {target_ip}, after {attempt} attempts, assuming connection to {expected_node_id} has dropped"
-                )
-                return None
-            time.sleep(1)
-            return _fetch_remote_node_id(attempt=attempt + 1)
-        except http.client.HTTPException as e:
-            logger.warning(f"HTTPException from {target_ip}: {type(e).__name__}: {e}")
-            return None
-        finally:
-            connection.close()
-
-    remote_node_id = await to_thread.run_sync(_fetch_remote_node_id)
-    if remote_node_id is None:
-        return
-
-    if remote_node_id == self_node_id:
-        return
-
-    if remote_node_id != expected_node_id:
-        logger.warning(
-            f"Discovered node with unexpected node_id; "
-            f"ip={target_ip}, expected_node_id={expected_node_id}, "
-            f"remote_node_id={remote_node_id}"
-        )
-        return
-
-    if remote_node_id not in out:
-        out[remote_node_id] = set()
-    out[remote_node_id].add(target_ip)
-
-
-async def check_reachable(
-    topology: Topology, self_node_id: NodeId
-) -> dict[NodeId, set[str]]:
-    """Check which nodes are reachable and return their IPs."""
-    reachable: dict[NodeId, set[str]] = {}
-    async with create_task_group() as tg:
-        for node in topology.list_nodes():
-            if not node.node_profile:
-                continue
-            for iface in node.node_profile.network_interfaces:
-                tg.start_soon(
-                    check_reachability,
-                    iface.ip_address,
-                    node.node_id,
-                    self_node_id,
-                    reachable,
-                )
-
-    return reachable

src/exo/worker/utils/profile.py

@@ -1,114 +0,0 @@
-import asyncio
-import os
-import platform
-from typing import Any, Callable, Coroutine
-
-import anyio
-from loguru import logger
-
-from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
-)
-
-from .macmon import (
-    MacMonError,
-    Metrics,
-)
-from .macmon import (
-    get_metrics_async as macmon_get_metrics_async,
-)
-from .system_info import (
-    get_friendly_name,
-    get_model_and_chip,
-    get_network_interfaces,
-)
-
-
-async def get_metrics_async() -> Metrics | None:
-    """Return detailed Metrics on macOS or a minimal fallback elsewhere."""
-
-    if platform.system().lower() == "darwin":
-        return await macmon_get_metrics_async()
-
-
-def get_memory_profile() -> MemoryPerformanceProfile:
-    """Construct a MemoryPerformanceProfile using psutil"""
-    override_memory_env = os.getenv("OVERRIDE_MEMORY_MB")
-    override_memory: int | None = (
-        Memory.from_mb(int(override_memory_env)).in_bytes
-        if override_memory_env
-        else None
-    )
-
-    return MemoryPerformanceProfile.from_psutil(override_memory=override_memory)
-
-
-async def start_polling_memory_metrics(
-    callback: Callable[[MemoryPerformanceProfile], Coroutine[Any, Any, None]],
-    *,
-    poll_interval_s: float = 0.5,
-) -> None:
-    """Continuously poll and emit memory-only metrics at a faster cadence.
-
-    Parameters
-    - callback: coroutine called with a fresh MemoryPerformanceProfile each tick
-    - poll_interval_s: interval between polls
-    """
-    while True:
-        try:
-            mem = get_memory_profile()
-            await callback(mem)
-        except MacMonError as e:
-            logger.opt(exception=e).error("Memory Monitor encountered error")
-        finally:
-            await anyio.sleep(poll_interval_s)
-
-
-async def start_polling_node_metrics(
-    callback: Callable[[NodePerformanceProfile], Coroutine[Any, Any, None]],
-):
-    poll_interval_s = 1.0
-    while True:
-        try:
-            metrics = await get_metrics_async()
-            if metrics is None:
-                return
-
-            network_interfaces = get_network_interfaces()
-            # these awaits could be joined but realistically they should be cached
-            model_id, chip_id = await get_model_and_chip()
-            friendly_name = await get_friendly_name()
-
-            # do the memory profile last to get a fresh reading to not conflict with the other memory profiling loop
-            memory_profile = get_memory_profile()
-
-            await callback(
-                NodePerformanceProfile(
-                    model_id=model_id,
-                    chip_id=chip_id,
-                    friendly_name=friendly_name,
-                    network_interfaces=network_interfaces,
-                    memory=memory_profile,
-                    system=SystemPerformanceProfile(
-                        gpu_usage=metrics.gpu_usage[1],
-                        temp=metrics.temp.gpu_temp_avg,
-                        sys_power=metrics.sys_power,
-                        pcpu_usage=metrics.pcpu_usage[1],
-                        ecpu_usage=metrics.ecpu_usage[1],
-                        ane_power=metrics.ane_power,
-                    ),
-                )
-            )
-
-        except asyncio.TimeoutError:
-            logger.warning(
-                "[resource_monitor] Operation timed out after 30s, skipping this cycle."
-            )
-        except MacMonError as e:
-            logger.opt(exception=e).error("Resource Monitor encountered error")
-            return
-        finally:
-            await anyio.sleep(poll_interval_s)

src/exo/worker/utils/tests/test_macmon.py

@@ -1,77 +0,0 @@
-"""Tests for macmon error handling.
-
-These tests verify that MacMon errors are handled gracefully without
-crashing the application or spamming logs.
-"""
-
-import platform
-from subprocess import CalledProcessError
-from unittest.mock import AsyncMock, patch
-
-import pytest
-
-from exo.worker.utils.macmon import MacMonError, get_metrics_async
-
-
-@pytest.mark.skipif(
-    platform.system().lower() != "darwin" or "arm" not in platform.machine().lower(),
-    reason="MacMon only supports macOS with Apple Silicon",
-)
-class TestMacMonErrorHandling:
-    """Test MacMon error handling."""
-
-    async def test_called_process_error_wrapped_as_macmon_error(self) -> None:
-        """CalledProcessError should be wrapped as MacMonError."""
-        mock_error = CalledProcessError(
-            returncode=1,
-            cmd=["macmon", "pipe", "-s", "1"],
-            stderr=b"some error message",
-        )
-
-        with (
-            patch(
-                "exo.worker.utils.macmon.shutil.which", return_value="/usr/bin/macmon"
-            ),
-            patch(
-                "exo.worker.utils.macmon.run_process", new_callable=AsyncMock
-            ) as mock_run,
-        ):
-            mock_run.side_effect = mock_error
-
-            with pytest.raises(MacMonError) as exc_info:
-                await get_metrics_async()
-
-            assert "MacMon failed with return code 1" in str(exc_info.value)
-            assert "some error message" in str(exc_info.value)
-
-    async def test_called_process_error_with_no_stderr(self) -> None:
-        """CalledProcessError with no stderr should be handled gracefully."""
-        mock_error = CalledProcessError(
-            returncode=1,
-            cmd=["macmon", "pipe", "-s", "1"],
-            stderr=None,
-        )
-
-        with (
-            patch(
-                "exo.worker.utils.macmon.shutil.which", return_value="/usr/bin/macmon"
-            ),
-            patch(
-                "exo.worker.utils.macmon.run_process", new_callable=AsyncMock
-            ) as mock_run,
-        ):
-            mock_run.side_effect = mock_error
-
-            with pytest.raises(MacMonError) as exc_info:
-                await get_metrics_async()
-
-            assert "MacMon failed with return code 1" in str(exc_info.value)
-            assert "no stderr" in str(exc_info.value)
-
-    async def test_macmon_not_found_raises_macmon_error(self) -> None:
-        """When macmon is not found in PATH, MacMonError should be raised."""
-        with patch("exo.worker.utils.macmon.shutil.which", return_value=None):
-            with pytest.raises(MacMonError) as exc_info:
-                await get_metrics_async()
-
-            assert "MacMon not found in PATH" in str(exc_info.value)

tests/headless_runner.py

@@ -34,7 +34,8 @@ from exo.shared.types.worker.instances import (
 )
 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, mp_channel
+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,
@@ -65,6 +66,7 @@ async def main():
     app = FastAPI()
     app.post("/ring")(ring_backend)
     app.post("/jaccl")(jaccl_backend)
+    app.post("/tb_detection")(tb_detection)
     shutdown = anyio.Event()
     await serve(
         app,  # type: ignore
@@ -76,6 +78,15 @@ async def main():
     shutdown.set()
 
 
+async def tb_detection():
+    send, recv = channel[GatheredInfo]()
+    ig = InfoGatherer(send)
+    with anyio.move_on_after(1):
+        await ig._monitor_system_profiler_thunderbolt_data()  # pyright: ignore[reportPrivateUsage]
+    with recv:
+        return recv.collect()
+
+
 async def assert_downloads():
     sd = exo_shard_downloader()
     # await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-0.6b"].model_id))
@@ -124,7 +135,7 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
     else:
         raise ValueError(f"{hn} not in {test.devs}")
 
-    meta = MODEL_CARDS[test.model_id].metadata
+    card = MODEL_CARDS[test.model_id]
     instance = MlxRingInstance(
         instance_id=iid,
         ephemeral_port=52416,
@@ -134,15 +145,15 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
             node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
             runner_to_shard={
                 RunnerId(test.devs[i][0]): PipelineShardMetadata(
-                    model_meta=meta,
+                    model_card=card,
                     device_rank=i,
                     world_size=world_size,
-                    start_layer=(meta.n_layers // world_size) * i,
+                    start_layer=(card.n_layers // world_size) * i,
                     end_layer=min(
-                        meta.n_layers, (meta.n_layers // world_size) * (i + 1)
+                        card.n_layers, (card.n_layers // world_size) * (i + 1)
                     ),
-                    n_layers=min(meta.n_layers, (meta.n_layers // world_size) * (i + 1))
-                    - (meta.n_layers // world_size) * i,
+                    n_layers=min(card.n_layers, (card.n_layers // world_size) * (i + 1))
+                    - (card.n_layers // world_size) * i,
                 )
                 for i in range(world_size)
             },
@@ -209,16 +220,16 @@ async def jaccl_backend(test: Tests):
             break
     else:
         raise ValueError(f"{weird_hn} not in {test.devs}")
-    return await execute_test(test, jaccl_instance(test, iid, hn), hn)
+    return await execute_test(test, jaccl_instance(test, iid), hn)
 
 
-def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
-    meta = MODEL_CARDS[test.model_id].metadata
+def jaccl_instance(test: Tests, iid: InstanceId):
+    card = MODEL_CARDS[test.model_id]
     world_size = len(test.devs)
 
     return MlxJacclInstance(
         instance_id=iid,
-        ibv_devices=[[None, "rdma_en3"], ["rdma_en3", None]],
+        jaccl_devices=[[None, "rdma_en3"], ["rdma_en3", None]],
         # rank 0 is always coordinator
         jaccl_coordinators={
             NodeId(host[0]): test.devs[0][1] + ":52416" for host in test.devs
@@ -228,12 +239,12 @@ def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
             node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
             runner_to_shard={
                 RunnerId(test.devs[i][0]): TensorShardMetadata(
-                    model_meta=meta,
+                    model_card=card,
                     device_rank=i,
                     world_size=world_size,
-                    start_layer=meta.n_layers,
-                    end_layer=meta.n_layers,
-                    n_layers=meta.n_layers,
+                    start_layer=card.n_layers,
+                    end_layer=card.n_layers,
+                    n_layers=card.n_layers,
                 )
                 for i in range(world_size)
             },

uv.lock

@@ -236,6 +236,7 @@ dependencies = [
     { name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+    { name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "hypercorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -268,11 +269,12 @@ requires-dist = [
     { name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
     { name = "fastapi", specifier = ">=0.116.1" },
     { name = "filelock", specifier = ">=3.18.0" },
+    { name = "httpx", specifier = ">=0.28.1" },
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
-    { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.1" },
-    { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.1" },
+    { 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", git = "https://github.com/AlexCheema/mlx-lm.git?rev=fix-transformers-5.0.0rc2" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "psutil", specifier = ">=7.0.0" },
@@ -633,22 +635,22 @@ wheels = [
 
 [[package]]
 name = "mlx"
-version = "0.30.1"
+version = "0.30.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "mlx-metal", marker = "sys_platform == 'darwin'" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/f9/fd/c6f56cd87d48763ed63655ace627c06db9819eae7d43d132f40d4965947a/mlx-0.30.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:743520758bc8261b2ed8f3b3dc96e4e9236769dd8f61fb17877c5e44037e2058", size = 593366 },
-    { url = "https://files.pythonhosted.org/packages/dc/53/96d8c48b21f91c4216b6d2ef6dfc10862e5fb0b811a2aaf02c96c78601de/mlx-0.30.1-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:fc9745bc1860ca60128e3a6d36157da06d936e2b4007a4dcba990b40202f598f", size = 593368 },
-    { url = "https://files.pythonhosted.org/packages/70/ce/476c3b7d3a4153bd0e1c5af1f1b6c09a804b652bbed34072404b322c22e0/mlx-0.30.1-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:a1480399c67bb327a66c5527b73915132e3fcaae3bce9634e5c81ccad9f43229", size = 567561 },
-    { url = "https://files.pythonhosted.org/packages/33/41/7ad1e639fd7dd1cf01a62c1c5b051024a859888c27504996e9d8380e6754/mlx-0.30.1-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:8e19850a4236a8e174f851f5789b8b62a8eb74f5a8fa49ad8ba286c5ddb5f9bf", size = 643122 },
-    { url = "https://files.pythonhosted.org/packages/d0/dc/72d3737c5b0662eb5e785d353dbc5e34d793d27b09b99e39993ee051bd19/mlx-0.30.1-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:1c8ed5bcd9f1910fca209e95859ac737e60b3e1954181b820fa269158f81049a", size = 687254 },
-    { url = "https://files.pythonhosted.org/packages/9b/cc/523448996247bb05d9d68e23bccf3dafdda660befb9330f6bd5fa13361e8/mlx-0.30.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d34cc2c25b0ee41c1349f14650db760e282685339858e305453f62405c12bc1b", size = 596006 },
-    { url = "https://files.pythonhosted.org/packages/23/0e/f9f2f9659c34c87be8f4167f6a1d6ed7e826f4889d20eecd4c0d8122f0e9/mlx-0.30.1-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:4e47d301e9095b87f0bda8827bfd6ffe744223aba5cee8f28e25894d647f5823", size = 596008 },
-    { url = "https://files.pythonhosted.org/packages/56/a7/49e41fb141de95b6a376091a963c737839c9cda04e423c67f57460a50458/mlx-0.30.1-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:cfba13e2a52255d663a1ad62f0f83eb3991e42147edf9a8d38cdd224e48ca49b", size = 570406 },
-    { url = "https://files.pythonhosted.org/packages/73/99/a43cb112167cf865c069f5e108ae42f5314663930ff3dd86c2d23d984191/mlx-0.30.1-cp314-cp314-manylinux_2_35_aarch64.whl", hash = "sha256:bebfec377208eb29cc88aa86c897c7446aa0984838669e138f273f9225d627ff", size = 646461 },
-    { url = "https://files.pythonhosted.org/packages/d4/ff/1e1968f107b4221a98dc26832586b1f646b27ddf3e55c95051c09d751f0a/mlx-0.30.1-cp314-cp314-manylinux_2_35_x86_64.whl", hash = "sha256:d18012d5cf0f013bc4a405cfd1e9d2d28e798f4d2dc4f15aa0fbffff73c02ba2", size = 687114 },
+    { url = "https://files.pythonhosted.org/packages/d0/22/42935d593fe82d3b98eb9d60e4620ed99703886635106f89d407c68f33bc/mlx-0.30.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:743fac1e4f9e8e46c8262943c643a31139c255cdb256c99ad496958215ccac1e", size = 569344 },
+    { url = "https://files.pythonhosted.org/packages/7d/27/f2e7a5236289d45315d0215e8553b4dd7e2faaba3bcb5025b34b25d5ab66/mlx-0.30.3-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:3b04ae81655aa0e63a6e8f2c749de3bbce64cf5b168ae10f39ed086dfa99e7f8", size = 569345 },
+    { url = "https://files.pythonhosted.org/packages/01/41/06b042457f51952456e9bb46b2c6e205ab3a28fc52d6751b5787fdb762b2/mlx-0.30.3-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:ba9b5bdb1e929cc130af72efd7f73508c0f4e526d224489af7ec1c6419564659", size = 569213 },
+    { url = "https://files.pythonhosted.org/packages/ec/1e/f62c98fc0d2d878ee4235671f9d406b13cc9240493ba6fcfde2f72c2ff83/mlx-0.30.3-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:dfe5c5b64e55398a22100804abbf9681996b03129e720e36b1727ed704db12b5", size = 617309 },
+    { url = "https://files.pythonhosted.org/packages/e9/62/811f064693449de740350d27793ce39343a460305ec8d878c318b80921d0/mlx-0.30.3-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:a3364924610929936e6aaf13c71106161258e5a5d3f7813a64c07cc2435f9f55", size = 659521 },
+    { url = "https://files.pythonhosted.org/packages/82/e2/6e551bd48fb350fbf0ee4cc5cd09485437d260b8f4937f22d8623e14687a/mlx-0.30.3-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:2c27fd8daaae14ca6cf407fcd236006a6e968f7708c8f61a2709116f2e754852", size = 571920 },
+    { url = "https://files.pythonhosted.org/packages/82/c0/561d1c9d3d12830b0e7fdcbd807585ef20909e398d4bcdbf25e4367543eb/mlx-0.30.3-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:b755fd4ed4b6a2ae4dee3766b5a2ea52fcbe83ebd1cf018458e18b74139409f3", size = 571921 },
+    { url = "https://files.pythonhosted.org/packages/42/1a/fb573fc2edc22a777fa254ff5c0c886ffd2c88aeb1f21c45778ef170f990/mlx-0.30.3-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:7e352c0369a2f7e54d4f317b434eab3333918ea9edde1c43c61d36386b6f76bf", size = 571732 },
+    { url = "https://files.pythonhosted.org/packages/9e/db/d0083e8f2205b3b2dcd9670eb6f0d6c1b7cbfea6b01a1f8bff39142edf44/mlx-0.30.3-cp314-cp314-manylinux_2_35_aarch64.whl", hash = "sha256:00ac867f3d003c1477a66a579442c2040ba7ea43ce3c174490d1f8bf379606bd", size = 619635 },
+    { url = "https://files.pythonhosted.org/packages/ab/90/ab0b93ff0e76da4fe0e878722c76a308cfb950b044a4676e9617276d8ccd/mlx-0.30.3-cp314-cp314-manylinux_2_35_x86_64.whl", hash = "sha256:5be7d0329036f09c6ed003ea3e307e97e3144f20a3e4711b01810d7d5013cf2c", size = 659652 },
 ]
 
 [package.optional-dependencies]
@@ -658,17 +660,17 @@ cpu = [
 
 [[package]]
 name = "mlx-cpu"
-version = "0.30.1"
+version = "0.30.3"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/64/51/32903727a68a61e972383e28a775c1f5e5f0628552c85cbc6103d68c0dc4/mlx_cpu-0.30.1-py3-none-manylinux_2_35_aarch64.whl", hash = "sha256:3f5dc2e4d0849181f8253508bb6a0854250483fc63d43ac79ec614b19824b172", size = 8992394 },
-    { url = "https://files.pythonhosted.org/packages/0c/74/69c21bb907f3c4064881ab0653029c939ae15fc4e63a5301ef8643cb1d68/mlx_cpu-0.30.1-py3-none-manylinux_2_35_x86_64.whl", hash = "sha256:c9ea6992d8c001e1123dfd3b4d4405ff576c787eec52656ad405e3d033a8be60", size = 10553055 },
+    { url = "https://files.pythonhosted.org/packages/91/8c/65bdb167141572ae0fa5d2da445368f454af52c0b93d8440b75df6472eef/mlx_cpu-0.30.3-py3-none-manylinux_2_35_aarch64.whl", hash = "sha256:85ba5601153e034d3ddb18b929147396737e74798c8453d7e078f53f35ae1512", size = 9000768 },
+    { url = "https://files.pythonhosted.org/packages/24/eb/e31b1edf9bb907add8fc2e386c6683004908acc71a0394e22856f139a9f1/mlx_cpu-0.30.3-py3-none-manylinux_2_35_x86_64.whl", hash = "sha256:5363249de88ab0f54af6ac00cf62d40c64f8e2ecd867598195d9baa9de8f6edb", size = 10569189 },
 ]
 
 [[package]]
 name = "mlx-lm"
-version = "0.30.2"
-source = { git = "https://github.com/AlexCheema/mlx-lm.git?rev=fix-transformers-5.0.0rc2#7180ee1f054fb6037c754fdf8259778f11300308" }
+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'" },
@@ -681,12 +683,12 @@ dependencies = [
 
 [[package]]
 name = "mlx-metal"
-version = "0.30.1"
+version = "0.30.3"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/09/3f/0be35ddad7e13d8ecd33a9185895f9739bb00b96ef0cce36cf0405d4aec0/mlx_metal-0.30.1-py3-none-macosx_14_0_arm64.whl", hash = "sha256:e7e92c6bdbd7ac8083f528a4c6640552ae106a57bb3d99856ac10a32e93a4b5e", size = 36864966 },
-    { url = "https://files.pythonhosted.org/packages/1e/1f/c0bddd0d5bf3871411aabe32121e09e1b7cdbece8917a49d5a442310e3e5/mlx_metal-0.30.1-py3-none-macosx_15_0_arm64.whl", hash = "sha256:bb50f57418af7fc3c42a2da2c4bde0e7ab7ac0b997de1f6f642a6680ac65d626", size = 36859011 },
-    { url = "https://files.pythonhosted.org/packages/67/b3/73cc2f584ac612a476096d35a61eed75ee7ed8b4e320b0c36cf60a14d4eb/mlx_metal-0.30.1-py3-none-macosx_26_0_arm64.whl", hash = "sha256:e0b151a0053ac00b4226710bfb6dbf54b87283fb01e10fb3877f9ea969f680aa", size = 44981160 },
+    { url = "https://files.pythonhosted.org/packages/f6/63/4d8f6fefb507c028df4454dabfe8d8e0ad2961bb06510b6aca23d2d5b2be/mlx_metal-0.30.3-py3-none-macosx_14_0_arm64.whl", hash = "sha256:6276312b02353714c7c6515169569fe1c4bebe3229c8ecf1fdb375a13e78c966", size = 37716245 },
+    { url = "https://files.pythonhosted.org/packages/35/91/1d452e48a4bb4958844fd3bb28ae31b8de110549c009ebec5024ce27ebf3/mlx_metal-0.30.3-py3-none-macosx_15_0_arm64.whl", hash = "sha256:c096c0a3428f3f96a06220f97a36f9528b18bc05173f821eb05bc8458e723fa8", size = 37712125 },
+    { url = "https://files.pythonhosted.org/packages/fe/36/7a3cbca85542b5ca4faf871e35927f43aa0e3fc830ae5b699780fe723677/mlx_metal-0.30.3-py3-none-macosx_26_0_arm64.whl", hash = "sha256:69068533bd1ee8b0379ce5de57ed5fd313577a10ecab58e1332fd1ff7248a75e", size = 46488962 },
 ]
 
 [[package]]