update log message + assertion

Dashboard fixes (TypeScript errors from `npm run check`):
- TopologyGraph.svelte: remove reference to deleted sendBackMultiaddr
  property, fix type inference for debug edge labels
- ModelCard.svelte: add missing topoWidth/topoHeight to early return
- +page.svelte: fix nested property access for deviceRank

Backend fix:
- info_gatherer.py: send initial MiscData on startup so friendly name
  appears immediately instead of showing "Unknown" until it changes

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

.github/workflows/build-app.yml

@@ -113,22 +113,11 @@ jobs:
           uv python install
           uv sync --locked
 
-      - name: Install Nix
-        uses: cachix/install-nix-action@v31
-        with:
-          nix_path: nixpkgs=channel:nixos-unstable
-
-      - name: Configure Cachix
-        uses: cachix/cachix-action@v14
-        with:
-          name: exo
-          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
-
       - name: Build dashboard
         run: |
-          DASHBOARD_OUT=$(nix build .#dashboard --print-build-logs --no-link --print-out-paths)
-          mkdir -p dashboard/build
-          cp -r "$DASHBOARD_OUT"/* dashboard/build/
+          cd dashboard
+          npm ci
+          npm run build
 
       - name: Install Sparkle CLI
         run: |

.github/workflows/pipeline.yml

@@ -20,12 +20,6 @@ jobs:
         with:
           nix_path: nixpkgs=channel:nixos-unstable
 
-      - uses: cachix/cachix-action@v14
-        name: Configure Cachix
-        with:
-          name: exo
-          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
-
       - name: Configure git user
         run: |
           git config --local user.email "github-actions@users.noreply.github.com"
@@ -94,19 +88,9 @@ jobs:
 
       - uses: ./.github/actions/typecheck
 
-  nix:
-    name: Build and check (${{ matrix.system }})
-    runs-on: ${{ matrix.runner }}
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - runner: macos-26
-            system: aarch64-darwin
-          - runner: ubuntu-latest
-            system: x86_64-linux
-          - runner: ubuntu-24.04-arm
-            system: aarch64-linux
+  nix-flake-check:
+    name: Check Nix flake
+    runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -117,20 +101,83 @@ jobs:
         with:
           nix_path: nixpkgs=channel:nixos-unstable
 
-      - uses: cachix/cachix-action@v14
-        name: Configure Cachix
-        with:
-          name: exo
-          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
-
-      - name: Build all Nix outputs
-        run: |
-          nix flake show --json | jq -r '
-            [
-              (.packages."${{ matrix.system }}" // {} | keys[] | ".#packages.${{ matrix.system }}.\(.)"),
-              (.devShells."${{ matrix.system }}" // {} | keys[] | ".#devShells.${{ matrix.system }}.\(.)")
-            ] | .[]
-          ' | xargs nix build
-
       - name: Run nix flake check
-        run: nix flake check
+        run: |
+          nix flake check
+        shell: bash
+
+#  ci:
+#    needs: typecheck
+#    runs-on: ubuntu-latest
+#    permissions:
+#      contents: read
+#    env:
+#      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+#    steps:
+#      - name: Checkout repository
+#        uses: actions/checkout@v4
+#        with:
+#          fetch-depth: 0
+#          token: ${{ secrets.GITHUB_TOKEN }}
+#          lfs: true
+#
+#      - name: Configure git user
+#        run: |
+#          git config --local user.email "github-actions@users.noreply.github.com"
+#          git config --local user.name  "github-actions bot"
+#        shell: bash
+#
+#      - name: Pull LFS files
+#        run: |
+#          echo "Pulling Git LFS files..."
+#          git lfs pull
+#        shell: bash
+#
+#      - name: Setup EXO_HOME and API_PORT
+#        run: |
+#          EXO_HOME=$(mktemp -d -t exo-ci-XXXXXXXX)
+#          # Generate random port (macOS compatible method)
+#          API_PORT=$((49152 + RANDOM % (65535 - 49152 + 1)))
+#          echo "EXO_HOME=$EXO_HOME" >> $GITHUB_ENV
+#          echo "API_PORT=$API_PORT" >> $GITHUB_ENV
+#          echo "Created EXO_HOME: $EXO_HOME"
+#          echo "Generated API_PORT: $API_PORT"
+#        shell: bash
+#
+#      - name: Setup Nix Environment
+#        run: |
+#          echo "Checking for nix installation..."
+#          
+#          # Check if nix binary exists directly
+#          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
+#            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
+#            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
+#            echo "PATH=$PATH" >> $GITHUB_ENV
+#            nix --version
+#          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
+#            echo "Found nix profile script, sourcing..."
+#            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
+#            nix --version
+#          elif command -v nix >/dev/null 2>&1; then
+#            echo "Nix already in PATH"
+#            nix --version
+#          else
+#            echo "Nix not found. Debugging info:"
+#            echo "Contents of /nix/var/nix/profiles/default/:"
+#            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
+#            echo "Contents of /nix/var/nix/profiles/default/bin/:"
+#            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
+#            exit 1
+#          fi
+#        shell: bash
+#
+#      - uses: ./.github/actions/lint-check
+#
+#      - uses: ./.github/actions/unit-test
+#
+#      - name: Cleanup EXO_HOME
+#        run: |
+#          echo "Cleaning up EXO_HOME: $EXO_HOME"
+#          rm -rf "$EXO_HOME"
+#        shell: bash
+#        if: always()

.mlx_typings/mlx_lm/models/deepseek_v3.pyi

@@ -1,156 +0,0 @@
-"""Type stubs for mlx_lm.models.deepseek_v3"""
-
-from dataclasses import dataclass
-from typing import Any, Dict, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs
-from .switch_layers import SwitchGLU
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    vocab_size: int
-    hidden_size: int
-    intermediate_size: int
-    moe_intermediate_size: int
-    num_hidden_layers: int
-    num_attention_heads: int
-    num_key_value_heads: int
-    n_shared_experts: Optional[int]
-    n_routed_experts: Optional[int]
-    routed_scaling_factor: float
-    kv_lora_rank: int
-    q_lora_rank: Optional[int]
-    qk_rope_head_dim: int
-    v_head_dim: int
-    qk_nope_head_dim: int
-    topk_method: str
-    scoring_func: str
-    norm_topk_prob: bool
-    n_group: int
-    topk_group: int
-    num_experts_per_tok: int
-    moe_layer_freq: int
-    first_k_dense_replace: int
-    max_position_embeddings: int
-    rms_norm_eps: float
-    rope_theta: float
-    rope_scaling: Optional[Dict[str, Any]]
-    attention_bias: bool
-
-class DeepseekV3Attention(nn.Module):
-    config: ModelArgs
-    hidden_size: int
-    num_heads: int
-    max_position_embeddings: int
-    rope_theta: float
-    q_lora_rank: Optional[int]
-    qk_rope_head_dim: int
-    kv_lora_rank: int
-    v_head_dim: int
-    qk_nope_head_dim: int
-    q_head_dim: int
-    scale: float
-    q_proj: nn.Linear
-    q_a_proj: nn.Linear
-    q_a_layernorm: nn.RMSNorm
-    q_b_proj: nn.Linear
-    kv_a_proj_with_mqa: nn.Linear
-    kv_a_layernorm: nn.RMSNorm
-    kv_b_proj: nn.Linear
-    o_proj: nn.Linear
-    rope: Any
-
-    def __init__(self, config: ModelArgs) -> None: ...
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array: ...
-
-class DeepseekV3MLP(nn.Module):
-    config: ModelArgs
-    hidden_size: int
-    intermediate_size: int
-    gate_proj: nn.Linear
-    up_proj: nn.Linear
-    down_proj: nn.Linear
-
-    def __init__(
-        self,
-        config: ModelArgs,
-        hidden_size: Optional[int] = None,
-        intermediate_size: Optional[int] = None,
-    ) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...
-
-class MoEGate(nn.Module):
-    config: ModelArgs
-    top_k: int
-    norm_topk_prob: bool
-    n_routed_experts: Optional[int]
-    routed_scaling_factor: float
-    n_group: int
-    topk_group: int
-    weight: mx.array
-    e_score_correction_bias: mx.array
-
-    def __init__(self, config: ModelArgs) -> None: ...
-    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
-
-class DeepseekV3MoE(nn.Module):
-    config: ModelArgs
-    num_experts_per_tok: int
-    switch_mlp: SwitchGLU
-    gate: MoEGate
-    shared_experts: DeepseekV3MLP
-    sharding_group: Optional[mx.distributed.Group]
-
-    def __init__(self, config: ModelArgs) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...
-
-class DeepseekV3DecoderLayer(nn.Module):
-    self_attn: DeepseekV3Attention
-    mlp: DeepseekV3MLP | DeepseekV3MoE
-    input_layernorm: nn.RMSNorm
-    post_attention_layernorm: nn.RMSNorm
-
-    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array: ...
-
-class DeepseekV3Model(nn.Module):
-    vocab_size: int
-    embed_tokens: nn.Embedding
-    layers: list[DeepseekV3DecoderLayer]
-    norm: nn.RMSNorm
-
-    def __init__(self, config: ModelArgs) -> None: ...
-    def __call__(
-        self,
-        x: mx.array,
-        cache: Optional[Any] = None,
-    ) -> mx.array: ...
-
-class Model(nn.Module):
-    model_type: str
-    model: DeepseekV3Model
-    lm_head: nn.Linear
-
-    def __init__(self, config: ModelArgs) -> None: ...
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[Any] = None,
-    ) -> mx.array: ...
-    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
-    @property
-    def layers(self) -> list[DeepseekV3DecoderLayer]: ...

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -57,11 +57,6 @@ class SwiGLU(nn.Module):
     def __call__(self, x, gate): ...
 
 class SwitchGLU(nn.Module):
-    gate_proj: SwitchLinear
-    up_proj: SwitchLinear
-    down_proj: SwitchLinear
-    activation: SwiGLU
-
     def __init__(
         self,
         input_dims: int,

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -4,7 +4,6 @@ This type stub file was generated by pyright.
 
 from functools import partial
 from pathlib import Path
-from typing import Any
 
 from transformers import PreTrainedTokenizerFast
 
@@ -104,55 +103,37 @@ class TokenizerWrapper:
     Accessing any attribute other than the ``detokenizer`` is forwarded to the
     huggingface tokenizer.
     """
+    def __init__(self, tokenizer, detokenizer_class=..., eos_token_ids=...) -> None: ...
+    def add_eos_token(self, token: str):  # -> None:
+        ...
+    @property
+    def has_thinking(self):  # -> bool:
+        ...
+    @property
+    def think_start(self):  # -> str | None:
+        ...
+    @property
+    def think_end(self):  # -> str | None:
+        ...
+    @property
+    def has_tool_calling(self):  # -> bool:
+        ...
+    @property
+    def tool_call_start(self):  # -> str | None:
+        ...
+    @property
+    def tool_call_end(self):  # -> str | None:
+        ...
+    @property
+    def detokenizer(self):  # -> NaiveStreamingDetokenizer:
+        """
+        Get a stateful streaming detokenizer.
+        """
 
-    _tokenizer: PreTrainedTokenizerFast
-    eos_token_id: int | None
-    eos_token: str | None
-    bos_token_id: int | None
-    bos_token: str | None
-    vocab_size: int
-    all_special_tokens: list[str]
-
-    def __init__(
-        self,
-        tokenizer: Any,
-        detokenizer_class: Any = ...,
-        eos_token_ids: list[int] | None = ...,
-        chat_template: Any = ...,
-        tool_parser: Any = ...,
-        tool_call_start: str | None = ...,
-        tool_call_end: str | None = ...,
-    ) -> None: ...
-    def encode(self, text: str, **kwargs: Any) -> list[int]: ...
-    def decode(self, token_ids: list[int], **kwargs: Any) -> str: ...
-    def apply_chat_template(
-        self,
-        messages: list[dict[str, Any]],
-        tokenize: bool = False,
-        add_generation_prompt: bool = False,
-        tools: Any = None,
-        **kwargs: Any,
-    ) -> str: ...
-    def get_vocab(self) -> dict[str, int]: ...
-    def add_eos_token(self, token: str) -> None: ...
-    @property
-    def has_thinking(self) -> bool: ...
-    @property
-    def think_start(self) -> str | None: ...
-    @property
-    def think_end(self) -> str | None: ...
-    @property
-    def has_tool_calling(self) -> bool: ...
-    @property
-    def tool_call_start(self) -> str | None: ...
-    @property
-    def tool_call_end(self) -> str | None: ...
-    @property
-    def detokenizer(self) -> NaiveStreamingDetokenizer:
-        """Get a stateful streaming detokenizer."""
-
-    def __getattr__(self, attr: str) -> Any: ...
-    def __setattr__(self, attr: str, value: Any) -> None: ...
+    def __getattr__(self, attr):  # -> set[Any] | Any:
+        ...
+    def __setattr__(self, attr, value):  # -> None:
+        ...
 
 class NewlineTokenizer(PreTrainedTokenizerFast):
     """A tokenizer that replaces newlines with <n> and <n> with new line."""
@@ -165,11 +146,18 @@ class NewlineTokenizer(PreTrainedTokenizerFast):
     def batch_decode(self, *args, **kwargs):  # -> list[str]:
         ...
 
-def load(
+def load_tokenizer(
     model_path: Path,
-    tokenizer_config_extra: dict[str, Any] | None = None,
-    eos_token_ids: list[int] | int | None = None,
-) -> TokenizerWrapper:
+    tokenizer_config_extra=...,
+    return_tokenizer=...,
+    eos_token_ids=...,
+) -> (
+    TokenizerWrapper
+    | type[SPMStreamingDetokenizer]
+    | partial[SPMStreamingDetokenizer]
+    | type[BPEStreamingDetokenizer]
+    | type[NaiveStreamingDetokenizer]
+):
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -177,7 +165,4 @@ def load(
     a Hugging Face repo ID.
     """
 
-# Alias for backward compatibility
-load_tokenizer = load
-
-def no_bos_or_eos(sequence: list[int], bos: int, eos: int) -> list[int]: ...
+def no_bos_or_eos(sequence: list, bos: int, eos: int) -> list: ...

AGENTS.md

@@ -1,103 +0,0 @@
-# AGENTS.md
-
-This file provides guidance to AI coding agents when working with code in this repository.
-
-## Project Overview
-
-exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and libp2p for peer-to-peer networking.
-
-## Build & Run Commands
-
-```bash
-# Build the dashboard (required before running exo)
-cd dashboard && npm install && npm run build && cd ..
-
-# Run exo (starts both master and worker with API at http://localhost:52415)
-uv run exo
-
-# Run with verbose logging
-uv run exo -v   # or -vv for more verbose
-
-# Run tests (excludes slow tests by default)
-uv run pytest
-
-# Run all tests including slow tests
-uv run pytest -m ""
-
-# Run a specific test file
-uv run pytest src/exo/shared/tests/test_election.py
-
-# Run a specific test function
-uv run pytest src/exo/shared/tests/test_election.py::test_function_name
-
-# Type checking (strict mode) - MUST pass before committing
-uv run basedpyright --project pyproject.toml
-
-# Linting
-uv run ruff check
-
-# Format code (using nix)
-nix fmt
-
-# Run all checks (do this before committing)
-uv run basedpyright --project pyproject.toml && uv run ruff check && nix fmt
-```
-
-## Architecture
-
-### Node Composition
-A single exo `Node` (src/exo/main.py) runs multiple components:
-- **Router**: libp2p-based pub/sub messaging via Rust bindings (exo_pyo3_bindings)
-- **Worker**: Handles inference tasks, downloads models, manages runner processes
-- **Master**: Coordinates cluster state, places model instances across nodes
-- **Election**: Bully algorithm for master election
-- **API**: FastAPI server for OpenAI-compatible chat completions
-
-### Message Flow
-Components communicate via typed pub/sub topics (src/exo/routing/topics.py):
-- `GLOBAL_EVENTS`: Master broadcasts indexed events to all workers
-- `LOCAL_EVENTS`: Workers send events to master for indexing
-- `COMMANDS`: Workers/API send commands to master
-- `ELECTION_MESSAGES`: Election protocol messages
-- `CONNECTION_MESSAGES`: libp2p connection updates
-
-### Event Sourcing
-The system uses event sourcing for state management:
-- `State` (src/exo/shared/types/state.py): Immutable state object
-- `apply()` (src/exo/shared/apply.py): Pure function that applies events to state
-- Master indexes events and broadcasts; workers apply indexed events
-
-### Key Type Hierarchy
-- `src/exo/shared/types/`: Pydantic models for all shared types
-  - `events.py`: Event types (discriminated union)
-  - `commands.py`: Command types
-  - `tasks.py`: Task types for worker execution
-  - `state.py`: Cluster state model
-
-### Rust Components
-Rust code in `rust/` provides:
-- `networking`: libp2p networking (gossipsub, peer discovery)
-- `exo_pyo3_bindings`: PyO3 bindings exposing Rust to Python
-- `system_custodian`: System-level operations
-
-### Dashboard
-Svelte 5 + TypeScript frontend in `dashboard/`. Build output goes to `dashboard/build/` and is served by the API.
-
-## Code Style Requirements
-
-From .cursorrules:
-- Strict, exhaustive typing - never bypass the type-checker
-- Use `Literal[...]` for enum-like sets, `typing.NewType` for primitives
-- Pydantic models with `frozen=True` and `strict=True`
-- Pure functions with injectable effect handlers for side-effects
-- Descriptive names - no abbreviations or 3-letter acronyms
-- Catch exceptions only where you can handle them meaningfully
-- Use `@final` and immutability wherever applicable
-
-## File Locations
-
-- **Downloaded models**: `~/.exo/models/` (NOT in huggingface cache)
-
-## 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.

CLAUDE.md

@@ -1 +0,0 @@
-AGENTS.md

Cargo.lock

@@ -4340,6 +4340,25 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "system_custodian"
+version = "0.0.1"
+dependencies = [
+ "delegate",
+ "derive_more",
+ "either",
+ "extend",
+ "futures",
+ "futures-timer",
+ "impl-trait-for-tuples",
+ "keccak-const",
+ "log",
+ "thiserror 2.0.17",
+ "tokio",
+ "tracing-subscriber",
+ "util",
+]
+
 [[package]]
 name = "tagptr"
 version = "0.2.0"

Cargo.toml

@@ -3,6 +3,7 @@ resolver = "3"
 members = [
     "rust/networking",
     "rust/exo_pyo3_bindings",
+    "rust/system_custodian",
     "rust/util",
 ]
 
@@ -24,6 +25,7 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
+system_custodian = { path = "rust/system_custodian" }
 util = { path = "rust/util" }
 
 # Proc-macro authoring tools

MISSED_THINGS.md

@@ -1,41 +0,0 @@
-# Missed things
-[X] Log EXO_LIBP2P_NAMESPACE on start in exo/main.py
-[X] Ordering of warmup was changed, which is wrong. It was changed to rank < n-1, then rank=n-1. It should be rank!=0 then rank=0 (this matches the auto_parallel implementation. NOTE: we use a different convention to mlx-lm, our terminal rank is rank=n-1 whereas mlx-lm is rank=0 hence i can see why this was changed wrongly).
-[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
-[X] Fetching download status of all models on start
-[X] Deduplication of tasks in plan_step.
-[X] resolve_allow_patterns should just be wildcard now.
-[] no mx_barrier in genreate.py mlx_generate at the end.
-[] cache assertion not needed in auto_parallel.py PipelineLastLayer.
-[] GPTOSS support dropped in auto_parallel.py.
-[] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
-[] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
-[] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
-[] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
-[X] KV_CACHE_BITS should be None to disable quantized KV cache.
-[] Dropped _set_nofile_limit in utils_mlx.py.
-[] We have group optional in load_mlx_items in utils_mlx.py.
-[] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
-[] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
-[X] We put cache limit back in utils_mlx.py.
-[] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
-[] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
-[] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
-[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
-[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
-[] logger.warning("You have likely selected ibv for a single node instance; falling back to MlxRing") was changed to debug. That will spam this warning since it happens every time we query instance previews.
-[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
-
-
-
-[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
-[X] Fetching download status of all models on start
-[X] Deduplication of tasks in plan_step.
-[X] resolve_allow_patterns should just be wildcard now.
-[X] KV_CACHE_BITS should be None to disable quantized KV cache.
-[X] We put cache limit back in utils_mlx.py.
-[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
-[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
-[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
-
-

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/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

dashboard/dashboard.nix

@@ -1,60 +0,0 @@
-{ lib
-, config
-, dream2nix
-, ...
-}:
-let
-  # Read and parse the lock file
-  rawLockFile = builtins.fromJSON (builtins.readFile "${config.deps.dashboardSrc}/package-lock.json");
-
-  # For packages with bundleDependencies, filter out deps that are bundled
-  # (bundled deps are inside the tarball, not separate lockfile entries)
-  fixedPackages = lib.mapAttrs
-    (path: entry:
-      if entry ? bundleDependencies && entry.bundleDependencies != [ ]
-      then entry // {
-        dependencies = lib.filterAttrs
-          (name: _: !(lib.elem name entry.bundleDependencies))
-          (entry.dependencies or { });
-      }
-      else entry
-    )
-    (rawLockFile.packages or { });
-
-  fixedLockFile = rawLockFile // { packages = fixedPackages; };
-in
-{
-  imports = [
-    dream2nix.modules.dream2nix.nodejs-package-lock-v3
-    dream2nix.modules.dream2nix.nodejs-granular-v3
-  ];
-
-  name = "exo-dashboard";
-  version = "1.0.0";
-
-  mkDerivation = {
-    src = config.deps.dashboardSrc;
-
-    buildPhase = ''
-      runHook preBuild
-      npm run build
-      runHook postBuild
-    '';
-
-    installPhase = ''
-      runHook preInstall
-      cp -r build $out/build
-      runHook postInstall
-    '';
-  };
-
-  deps = { nixpkgs, ... }: {
-    inherit (nixpkgs) stdenv;
-    dashboardSrc = null; # Injected by parts.nix
-  };
-
-  nodejs-package-lock-v3 = {
-    # Don't use packageLockFile - provide the fixed lock content directly
-    packageLock = fixedLockFile;
-  };
-}

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/parts.nix

@@ -1,44 +0,0 @@
-{ inputs, ... }:
-{
-  perSystem =
-    { pkgs, lib, ... }:
-    let
-      # Filter source to only include dashboard directory
-      src = lib.cleanSourceWith {
-        src = inputs.self;
-        filter =
-          path: type:
-          let
-            baseName = builtins.baseNameOf path;
-            inDashboardDir =
-              (lib.hasInfix "/dashboard/" path)
-              || (lib.hasSuffix "/dashboard" (builtins.dirOf path))
-              || (baseName == "dashboard" && type == "directory");
-          in
-          inDashboardDir;
-      };
-
-      # Build the dashboard with dream2nix (includes node_modules in output)
-      dashboardFull = inputs.dream2nix.lib.evalModules {
-        packageSets.nixpkgs = pkgs;
-        modules = [
-          ./dashboard.nix
-          {
-            paths.projectRoot = inputs.self;
-            paths.projectRootFile = "flake.nix";
-            paths.package = inputs.self + "/dashboard";
-          }
-          # Inject the filtered source
-          {
-            deps.dashboardSrc = lib.mkForce "${src}/dashboard";
-          }
-        ];
-      };
-    in
-    {
-      # Extract just the static site from the full build
-      packages.dashboard = pkgs.runCommand "exo-dashboard" { } ''
-        cp -r ${dashboardFull}/build $out
-      '';
-    };
-}

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

@@ -99,7 +99,7 @@ interface RawNodeProfile {
 
 interface RawTopologyNode {
 	nodeId: string;
-	nodeProfile: RawNodeProfile;
+	nodeProfile?: RawNodeProfile;
 }
 
 interface RawTopologyConnection {
@@ -110,9 +110,19 @@ interface RawTopologyConnection {
 		| string;
 }
 
+// Connection can be an object or a tuple [source, target, metadata]
+type RawConnectionItem =
+	| RawTopologyConnection
+	| [
+			string,
+			string,
+			{ sinkMultiaddr?: { ip_address?: string; address?: string } }?,
+	  ];
+
 interface RawTopology {
-	nodes: RawTopologyNode[];
-	connections?: RawTopologyConnection[];
+	// nodes can be array of strings (node IDs) or array of objects with nodeId/nodeProfile
+	nodes: (string | RawTopologyNode)[];
+	connections?: RawConnectionItem[];
 }
 
 type RawNodeProfiles = Record<string, RawNodeProfile>;
@@ -213,9 +223,18 @@ function transformTopology(
 	const nodes: Record<string, NodeInfo> = {};
 	const edges: TopologyEdge[] = [];
 
+	// Handle nodes - can be array of strings (node IDs) or array of objects with nodeId/nodeProfile
 	for (const node of raw.nodes || []) {
-		const mergedProfile = profiles?.[node.nodeId];
-		const profile = { ...(node.nodeProfile ?? {}), ...(mergedProfile ?? {}) };
+		// Determine the node ID - could be a string or an object with nodeId property
+		const nodeId = typeof node === "string" ? node : node.nodeId;
+		if (!nodeId) continue;
+
+		// Get the profile - from the separate profiles map or from the node object itself
+		const profileFromMap = profiles?.[nodeId];
+		const profileFromNode =
+			typeof node === "object" ? node.nodeProfile : undefined;
+		const profile = { ...(profileFromNode ?? {}), ...(profileFromMap ?? {}) };
+
 		const ramTotal = profile?.memory?.ramTotal?.inBytes ?? 0;
 		const ramAvailable = profile?.memory?.ramAvailable?.inBytes ?? 0;
 		const ramUsage = Math.max(ramTotal - ramAvailable, 0);
@@ -264,7 +283,7 @@ function transformTopology(
 			}
 		}
 
-		nodes[node.nodeId] = {
+		nodes[nodeId] = {
 			system_info: {
 				model_id: profile?.modelId ?? "Unknown",
 				chip: profile?.chipId,
@@ -292,14 +311,39 @@ function transformTopology(
 		};
 	}
 
+	// Handle connections - can be objects with localNodeId/sendBackNodeId or tuples [source, target, metadata]
 	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;
+		let localNodeId: string | undefined;
+		let sendBackNodeId: string | undefined;
+		let sendBackMultiaddr:
+			| { multiaddr?: string; address?: string; ip_address?: string }
+			| string
+			| undefined;
+
+		// Check if it's a tuple format [source, target, metadata]
+		if (Array.isArray(conn)) {
+			localNodeId = conn[0] as string;
+			sendBackNodeId = conn[1] as string;
+			const metadata = conn[2] as
+				| { sinkMultiaddr?: { ip_address?: string; address?: string } }
+				| undefined;
+			if (metadata?.sinkMultiaddr) {
+				sendBackMultiaddr = metadata.sinkMultiaddr;
+			}
+		} else {
+			// Object format with localNodeId/sendBackNodeId
+			localNodeId = conn.localNodeId;
+			sendBackNodeId = conn.sendBackNodeId;
+			sendBackMultiaddr = conn.sendBackMultiaddr;
+		}
+
+		if (!localNodeId || !sendBackNodeId) continue;
+		if (localNodeId === sendBackNodeId) continue;
+		if (!nodes[localNodeId] || !nodes[sendBackNodeId]) continue;
 
 		let sendBackIp: string | undefined;
-		if (conn.sendBackMultiaddr) {
-			const multi = conn.sendBackMultiaddr;
+		if (sendBackMultiaddr) {
+			const multi = sendBackMultiaddr;
 			if (typeof multi === "string") {
 				sendBackIp = extractIpFromMultiaddr(multi);
 			} else {
@@ -311,8 +355,8 @@ function transformTopology(
 		}
 
 		edges.push({
-			source: conn.localNodeId,
-			target: conn.sendBackNodeId,
+			source: localNodeId,
+			target: sendBackNodeId,
 			sendBackIp,
 		});
 	}

dashboard/src/routes/+page.svelte

@@ -895,7 +895,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

@@ -199,13 +199,7 @@
 					const rawProgress = (downloadPayload as Record<string, unknown>).download_progress
 						?? (downloadPayload as Record<string, unknown>).downloadProgress
 						?? {};
-					// For DownloadCompleted, total_bytes is at top level; for DownloadOngoing, it's inside download_progress
-					const totalBytes = getBytes(
-						(downloadPayload as Record<string, unknown>).total_bytes
-						?? (downloadPayload as Record<string, unknown>).totalBytes
-						?? (rawProgress as Record<string, unknown>).total_bytes
-						?? (rawProgress as Record<string, unknown>).totalBytes
-					);
+					const totalBytes = getBytes((rawProgress as Record<string, unknown>).total_bytes ?? (rawProgress as Record<string, unknown>).totalBytes);
 					const downloadedBytes = getBytes((rawProgress as Record<string, unknown>).downloaded_bytes ?? (rawProgress as Record<string, unknown>).downloadedBytes);
 					const speed = (rawProgress as Record<string, unknown>).speed as number ?? 0;
 					const etaMs = (rawProgress as Record<string, unknown>).eta_ms as number ?? (rawProgress as Record<string, unknown>).etaMs as number ?? 0;
@@ -338,13 +332,8 @@
 								<div class="text-lg font-mono text-white truncate">{node.nodeName}</div>
 								<div class="text-xs text-exo-light-gray font-mono truncate">{node.nodeId}</div>
 							</div>
-							<div class="text-xs font-mono uppercase tracking-wider whitespace-nowrap shrink-0 text-right">
-								<div>
-									<span class="text-green-400">{node.models.filter(m => m.status === 'completed').length}</span><span class="text-exo-yellow"> / {node.models.length} models</span>
-								</div>
-								<div class="text-exo-light-gray normal-case tracking-normal">
-									{formatBytes(node.models.filter(m => m.status === 'completed').reduce((sum, m) => sum + m.totalBytes, 0))} on disk
-								</div>
+							<div class="text-xs font-mono uppercase tracking-wider whitespace-nowrap shrink-0">
+								<span class="text-green-400">{node.models.filter(m => m.status === 'completed').length}</span><span class="text-exo-yellow"> /{node.models.length} models</span>
 							</div>
 						</div>
 
@@ -396,7 +385,7 @@
 								</div>
 
 								<div class="flex items-center justify-between text-xs font-mono text-exo-light-gray">
-									<span>{model.status === 'completed' ? `Completed (${formatBytes(model.totalBytes)})` : `${formatSpeed(model.speed)} • ETA ${formatEta(model.etaMs)}`}</span>
+									<span>{model.status === 'completed' ? 'Completed' : `${formatSpeed(model.speed)} • ETA ${formatEta(model.etaMs)}`}</span>
 									{#if model.status !== 'completed'}
 										<span>{model.files.length} file{model.files.length === 1 ? '' : 's'}</span>
 									{/if}

flake.lock

@@ -1,42 +1,5 @@
 {
   "nodes": {
-    "crane": {
-      "locked": {
-        "lastModified": 1767744144,
-        "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=",
-        "owner": "ipetkov",
-        "repo": "crane",
-        "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "ipetkov",
-        "repo": "crane",
-        "type": "github"
-      }
-    },
-    "dream2nix": {
-      "inputs": {
-        "nixpkgs": [
-          "nixpkgs"
-        ],
-        "purescript-overlay": "purescript-overlay",
-        "pyproject-nix": "pyproject-nix"
-      },
-      "locked": {
-        "lastModified": 1765953015,
-        "narHash": "sha256-5FBZbbWR1Csp3Y2icfRkxMJw/a/5FGg8hCXej2//bbI=",
-        "owner": "nix-community",
-        "repo": "dream2nix",
-        "rev": "69eb01fa0995e1e90add49d8ca5bcba213b0416f",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-community",
-        "repo": "dream2nix",
-        "type": "github"
-      }
-    },
     "fenix": {
       "inputs": {
         "nixpkgs": [
@@ -45,11 +8,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1768287139,
-        "narHash": "sha256-nsXFt0OzUi6K7dUzzJD5/v9e0Ic+fvclfIW936/43ZM=",
+        "lastModified": 1761893049,
+        "narHash": "sha256-1TtFDPhC+ZsrOOtBnry1EZC+WipTTvsOVjIEVugqji8=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "a4a3aa956931f90f35453cb519e4545e9ad7f773",
+        "rev": "c2ac9a5c0d6d16630c3b225b874bd14528d1abe6",
         "type": "github"
       },
       "original": {
@@ -58,59 +21,25 @@
         "type": "github"
       }
     },
-    "flake-compat": {
-      "flake": false,
-      "locked": {
-        "lastModified": 1696426674,
-        "narHash": "sha256-kvjfFW7WAETZlt09AgDn1MrtKzP7t90Vf7vypd3OL1U=",
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "rev": "0f9255e01c2351cc7d116c072cb317785dd33b33",
-        "type": "github"
-      },
-      "original": {
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "type": "github"
-      }
-    },
-    "flake-parts": {
+    "flake-utils": {
       "inputs": {
-        "nixpkgs-lib": [
-          "nixpkgs"
-        ]
+        "systems": "systems"
       },
       "locked": {
-        "lastModified": 1768135262,
-        "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=",
-        "owner": "hercules-ci",
-        "repo": "flake-parts",
-        "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac",
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
         "type": "github"
       },
       "original": {
-        "owner": "hercules-ci",
-        "repo": "flake-parts",
+        "owner": "numtide",
+        "repo": "flake-utils",
         "type": "github"
       }
     },
     "nixpkgs": {
-      "locked": {
-        "lastModified": 1768127708,
-        "narHash": "sha256-1Sm77VfZh3mU0F5OqKABNLWxOuDeHIlcFjsXeeiPazs=",
-        "owner": "NixOS",
-        "repo": "nixpkgs",
-        "rev": "ffbc9f8cbaacfb331b6017d5a5abb21a492c9a38",
-        "type": "github"
-      },
-      "original": {
-        "owner": "NixOS",
-        "ref": "nixos-unstable",
-        "repo": "nixpkgs",
-        "type": "github"
-      }
-    },
-    "nixpkgs-swift": {
       "locked": {
         "lastModified": 1761672384,
         "narHash": "sha256-o9KF3DJL7g7iYMZq9SWgfS1BFlNbsm6xplRjVlOCkXI=",
@@ -121,74 +50,27 @@
       },
       "original": {
         "owner": "NixOS",
+        "ref": "nixos-unstable",
         "repo": "nixpkgs",
-        "rev": "08dacfca559e1d7da38f3cf05f1f45ee9bfd213c",
-        "type": "github"
-      }
-    },
-    "purescript-overlay": {
-      "inputs": {
-        "flake-compat": "flake-compat",
-        "nixpkgs": [
-          "dream2nix",
-          "nixpkgs"
-        ],
-        "slimlock": "slimlock"
-      },
-      "locked": {
-        "lastModified": 1728546539,
-        "narHash": "sha256-Sws7w0tlnjD+Bjck1nv29NjC5DbL6nH5auL9Ex9Iz2A=",
-        "owner": "thomashoneyman",
-        "repo": "purescript-overlay",
-        "rev": "4ad4c15d07bd899d7346b331f377606631eb0ee4",
-        "type": "github"
-      },
-      "original": {
-        "owner": "thomashoneyman",
-        "repo": "purescript-overlay",
-        "type": "github"
-      }
-    },
-    "pyproject-nix": {
-      "inputs": {
-        "nixpkgs": [
-          "dream2nix",
-          "nixpkgs"
-        ]
-      },
-      "locked": {
-        "lastModified": 1763017646,
-        "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=",
-        "owner": "pyproject-nix",
-        "repo": "pyproject.nix",
-        "rev": "47bd6f296502842643078d66128f7b5e5370790c",
-        "type": "github"
-      },
-      "original": {
-        "owner": "pyproject-nix",
-        "repo": "pyproject.nix",
         "type": "github"
       }
     },
     "root": {
       "inputs": {
-        "crane": "crane",
-        "dream2nix": "dream2nix",
         "fenix": "fenix",
-        "flake-parts": "flake-parts",
+        "flake-utils": "flake-utils",
         "nixpkgs": "nixpkgs",
-        "nixpkgs-swift": "nixpkgs-swift",
         "treefmt-nix": "treefmt-nix"
       }
     },
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1768224240,
-        "narHash": "sha256-Pp1dDrXKPBUJReZnnDElFyHYn67XTd48zRhToheLjtk=",
+        "lastModified": 1761849405,
+        "narHash": "sha256-igXdvC+WCUN+3gnfk+ptT7rMmxQuY6WbIg1rXMUN1DM=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "725349602e525df37f377701e001fe8aab807878",
+        "rev": "f7de8ae045a5fe80f1203c5a1c3015b05f7c3550",
         "type": "github"
       },
       "original": {
@@ -198,25 +80,18 @@
         "type": "github"
       }
     },
-    "slimlock": {
-      "inputs": {
-        "nixpkgs": [
-          "dream2nix",
-          "purescript-overlay",
-          "nixpkgs"
-        ]
-      },
+    "systems": {
       "locked": {
-        "lastModified": 1688756706,
-        "narHash": "sha256-xzkkMv3neJJJ89zo3o2ojp7nFeaZc2G0fYwNXNJRFlo=",
-        "owner": "thomashoneyman",
-        "repo": "slimlock",
-        "rev": "cf72723f59e2340d24881fd7bf61cb113b4c407c",
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
         "type": "github"
       },
       "original": {
-        "owner": "thomashoneyman",
-        "repo": "slimlock",
+        "owner": "nix-systems",
+        "repo": "default",
         "type": "github"
       }
     },
@@ -227,11 +102,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768158989,
-        "narHash": "sha256-67vyT1+xClLldnumAzCTBvU0jLZ1YBcf4vANRWP3+Ak=",
+        "lastModified": 1762938485,
+        "narHash": "sha256-AlEObg0syDl+Spi4LsZIBrjw+snSVU4T8MOeuZJUJjM=",
         "owner": "numtide",
         "repo": "treefmt-nix",
-        "rev": "e96d59dff5c0d7fddb9d113ba108f03c3ef99eca",
+        "rev": "5b4ee75aeefd1e2d5a1cc43cf6ba65eba75e83e4",
         "type": "github"
       },
       "original": {

flake.nix

@@ -3,134 +3,132 @@
 
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
-
-    flake-parts = {
-      url = "github:hercules-ci/flake-parts";
-      inputs.nixpkgs-lib.follows = "nixpkgs";
-    };
-
-    crane.url = "github:ipetkov/crane";
-
+    flake-utils.url = "github:numtide/flake-utils";
+    # Provides Rust dev-env integration:
     fenix = {
       url = "github:nix-community/fenix";
       inputs.nixpkgs.follows = "nixpkgs";
     };
-
+    # Provides formatting infrastructure:
     treefmt-nix = {
       url = "github:numtide/treefmt-nix";
       inputs.nixpkgs.follows = "nixpkgs";
     };
-
-    dream2nix = {
-      url = "github:nix-community/dream2nix";
-      inputs.nixpkgs.follows = "nixpkgs";
-    };
-
-    # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
-    nixpkgs-swift.url = "github:NixOS/nixpkgs/08dacfca559e1d7da38f3cf05f1f45ee9bfd213c";
   };
 
-  nixConfig = {
-    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI=";
-    extra-substituters = "https://exo.cachix.org";
-  };
+  # TODO: figure out caching story
+  # nixConfig = {
+  #   # nix community cachix
+  #   extra-trusted-public-keys = "nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs=";
+  #   extra-substituters = "https://nix-community.cachix.org";
+  # };
 
   outputs =
     inputs:
-    inputs.flake-parts.lib.mkFlake { inherit inputs; } {
+    let
       systems = [
         "x86_64-linux"
         "aarch64-darwin"
         "aarch64-linux"
       ];
-
-      imports = [
-        inputs.treefmt-nix.flakeModule
-        ./dashboard/parts.nix
-        ./rust/parts.nix
-      ];
-
-      perSystem =
-        { config, self', inputs', pkgs, lib, system, ... }:
-        let
-          fenixToolchain = inputs'.fenix.packages.complete;
-          # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
-          pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
-        in
-        {
-          treefmt = {
-            projectRootFile = "flake.nix";
-            programs = {
-              nixpkgs-fmt.enable = true;
-              ruff-format = {
-                enable = true;
-                excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
-              };
-              rustfmt = {
-                enable = true;
-                package = config.rust.toolchain;
-              };
-              prettier = {
-                enable = true;
-                includes = [ "*.ts" ];
-              };
-              swift-format = {
-                enable = true;
-                package = pkgsSwift.swiftPackages.swift-format;
-              };
+      fenixToolchain = system: inputs.fenix.packages.${system}.complete;
+    in
+    inputs.flake-utils.lib.eachSystem systems (
+      system:
+      let
+        pkgs = import inputs.nixpkgs {
+          inherit system;
+          overlays = [ inputs.fenix.overlays.default ];
+        };
+        treefmtEval = inputs.treefmt-nix.lib.evalModule pkgs {
+          projectRootFile = "flake.nix";
+          programs = {
+            nixpkgs-fmt.enable = true;
+            ruff-format = {
+              enable = true;
+              excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
             };
-          };
-
-          checks.lint = pkgs.runCommand "lint-check" { } ''
-            export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
-            ${pkgs.ruff}/bin/ruff check ${inputs.self}/
-            touch $out
-          '';
-
-          devShells.default = with pkgs; pkgs.mkShell {
-            inputsFrom = [ self'.checks.cargo-build ];
-
-            packages =
-              [
-                # FORMATTING
-                config.treefmt.build.wrapper
-
-                # PYTHON
-                python313
-                uv
-                ruff
-                basedpyright
-
-                # RUST
-                config.rust.toolchain
-                maturin
-
-                # NIX
-                nixpkgs-fmt
-
-                # SVELTE
-                nodejs
-
-                # MISC
-                just
-                jq
-              ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
-              ++ lib.optionals stdenv.isDarwin [
-                macmon
-              ];
-
-            OPENSSL_NO_VENDOR = "1";
-
-            shellHook = ''
-              export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
-              ${lib.optionalString stdenv.isLinux ''
-                export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
-              ''}
-            '';
+            rustfmt = {
+              enable = true;
+              package = (fenixToolchain system).rustfmt;
+            };
+            prettier = {
+              enable = true;
+              includes = [ "*.ts" ];
+            };
+            swift-format.enable = true;
           };
         };
-    };
+      in
+      {
+        formatter = treefmtEval.config.build.wrapper;
+        checks.formatting = treefmtEval.config.build.check inputs.self;
+        checks.lint = pkgs.runCommand "lint-check" { } ''
+          export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+          ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+          touch $out
+        '';
+
+        devShells.default = pkgs.mkShell {
+          packages =
+            with pkgs;
+            [
+              # FORMATTING
+              treefmtEval.config.build.wrapper
+
+              # PYTHON
+              python313
+              uv
+              ruff
+              basedpyright
+
+              # RUST
+              ((fenixToolchain system).withComponents [
+                "cargo"
+                "rustc"
+                "clippy"
+                "rustfmt"
+                "rust-src"
+              ])
+              rustup # Just here to make RustRover happy
+
+              # NIX
+              nixpkgs-fmt
+
+              # SVELTE
+              nodejs
+
+              # MISC
+              just
+              jq
+            ]
+            ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [
+              # IFCONFIG
+              unixtools.ifconfig
+
+              # Build dependencies for Linux
+              pkg-config
+              openssl
+            ])
+            ++ (pkgs.lib.optionals pkgs.stdenv.isDarwin [
+              # MACMON
+              macmon
+            ]);
+
+          shellHook = ''
+            # PYTHON
+            export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${pkgs.python313}/lib"
+            ${pkgs.lib.optionalString pkgs.stdenv.isLinux ''
+              # Build environment for Linux
+              export PKG_CONFIG_PATH="${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH"
+              export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH"
+            ''}
+            echo
+            echo "🍎🍎 Run 'just <recipe>' to get started"
+            just --list
+          '';
+
+        };
+      }
+    );
 }

pyproject.toml

@@ -17,9 +17,9 @@ 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-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
+    "mlx>=0.30.1; sys_platform == 'darwin'",
+    "mlx[cpu]>=0.30.1; sys_platform == 'linux'",
+    "mlx-lm>=0.28.3",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
@@ -33,7 +33,6 @@ exo = "exo.main:main"
 # dependencies only required for development
 [dependency-groups]
 dev = [
-    "basedpyright>=1.29.0",
     "pyinstaller>=6.17.0",
     "pytest>=8.4.0",
     "pytest-asyncio>=1.0.0",
@@ -99,7 +98,6 @@ root = "src"
 
 # supported platforms for this project
 [tool.uv]
-prerelease = "allow"
 environments = [
     "sys_platform == 'darwin'",
     "sys_platform == 'linux'",

rust/parts.nix

@@ -1,145 +0,0 @@
-{ inputs, ... }:
-{
-  perSystem =
-    { config, self', inputs', pkgs, lib, ... }:
-    let
-      # Fenix nightly toolchain with all components
-      fenixPkgs = inputs'.fenix.packages;
-      rustToolchain = fenixPkgs.complete.withComponents [
-        "cargo"
-        "rustc"
-        "clippy"
-        "rustfmt"
-        "rust-src"
-        "rust-analyzer"
-      ];
-
-      # Crane with fenix toolchain
-      craneLib = (inputs.crane.mkLib pkgs).overrideToolchain rustToolchain;
-
-      # Source filtering - only include rust/ directory and root Cargo files
-      # This ensures changes to Python/docs/etc don't trigger Rust rebuilds
-      src = lib.cleanSourceWith {
-        src = inputs.self;
-        filter =
-          path: type:
-          let
-            baseName = builtins.baseNameOf path;
-            parentDir = builtins.dirOf path;
-            inRustDir =
-              (lib.hasInfix "/rust/" path)
-              || (lib.hasSuffix "/rust" parentDir)
-              || (baseName == "rust" && type == "directory");
-            isRootCargoFile =
-              (baseName == "Cargo.toml" || baseName == "Cargo.lock")
-              && (builtins.dirOf path == toString inputs.self);
-          in
-          isRootCargoFile
-          || (inRustDir && (craneLib.filterCargoSources path type || lib.hasSuffix ".toml" path || lib.hasSuffix ".md" path));
-      };
-
-      # Common arguments for all Rust builds
-      commonArgs = {
-        inherit src;
-        pname = "exo-rust";
-        version = "0.0.1";
-        strictDeps = true;
-
-        nativeBuildInputs = [
-          pkgs.pkg-config
-          pkgs.python313 # Required for pyo3-build-config
-        ];
-
-        buildInputs = [
-          pkgs.openssl
-          pkgs.python313 # Required for pyo3 tests
-        ];
-
-        OPENSSL_NO_VENDOR = "1";
-
-        # Required for pyo3 tests to find libpython
-        LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ];
-      };
-
-      # Build dependencies once for caching
-      cargoArtifacts = craneLib.buildDepsOnly (
-        commonArgs
-        // {
-          cargoExtraArgs = "--workspace";
-        }
-      );
-    in
-    {
-      # Export toolchain for use in treefmt and devShell
-      options.rust = {
-        toolchain = lib.mkOption {
-          type = lib.types.package;
-          default = rustToolchain;
-          description = "The Rust toolchain to use";
-        };
-      };
-
-      config = {
-        packages = {
-          # Python bindings wheel via maturin
-          exo_pyo3_bindings = craneLib.buildPackage (
-            commonArgs
-            // {
-              inherit cargoArtifacts;
-              pname = "exo_pyo3_bindings";
-
-              nativeBuildInputs = commonArgs.nativeBuildInputs ++ [
-                pkgs.maturin
-              ];
-
-              buildPhaseCargoCommand = ''
-                maturin build \
-                  --release \
-                  --manylinux off \
-                  --manifest-path rust/exo_pyo3_bindings/Cargo.toml \
-                  --features "pyo3/extension-module,pyo3/experimental-async" \
-                  --interpreter ${pkgs.python313}/bin/python \
-                  --out dist
-              '';
-
-              # Don't use crane's default install behavior
-              doNotPostBuildInstallCargoBinaries = true;
-
-              installPhaseCommand = ''
-                mkdir -p $out
-                cp dist/*.whl $out/
-              '';
-            }
-          );
-        };
-
-        checks = {
-          # Full workspace build (all crates)
-          cargo-build = craneLib.buildPackage (
-            commonArgs
-            // {
-              inherit cargoArtifacts;
-              cargoExtraArgs = "--workspace";
-            }
-          );
-          # Run tests with nextest
-          cargo-nextest = craneLib.cargoNextest (
-            commonArgs
-            // {
-              inherit cargoArtifacts;
-              cargoExtraArgs = "--workspace";
-            }
-          );
-
-          # Build documentation
-          cargo-doc = craneLib.cargoDoc (
-            commonArgs
-            // {
-              inherit cargoArtifacts;
-              cargoExtraArgs = "--workspace";
-            }
-          );
-        };
-      };
-    };
-}

rust/system_custodian/Cargo.toml

@@ -0,0 +1,47 @@
+[package]
+name = "system_custodian"
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
+
+[lib]
+doctest = false
+name = "system_custodian"
+path = "src/lib.rs"
+
+[[bin]]
+path = "src/bin/main.rs"
+name = "system_custodian"
+doc = false
+
+[lints]
+workspace = true
+
+[dependencies]
+# datastructures
+either = { workspace = true }
+
+# macro dependencies
+extend = { workspace = true }
+delegate = { workspace = true }
+impl-trait-for-tuples = { workspace = true }
+derive_more = { workspace = true }
+
+# async
+tokio = { workspace = true, features = ["full"] }
+futures = { workspace = true }
+futures-timer = { workspace = true }
+
+# utility dependencies
+util = { workspace = true }
+thiserror = { workspace = true }
+#internment = { workspace = true }
+#recursion = { workspace = true }
+#generativity = { workspace = true }
+#itertools = { workspace = true }
+tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
+keccak-const = { workspace = true }
+
+# tracing/logging
+log = { workspace = true }
+

rust/system_custodian/src/bin/main.rs

@@ -0,0 +1,4 @@
+//! TODO: documentation
+//!
+
+fn main() {}

rust/system_custodian/src/lib.rs

@@ -0,0 +1,69 @@
+//! This crate defines the logic of, and ways to interact with, Exo's **_System Custodian_** daemon.
+//!
+//! The **_System Custodian_** daemon is supposed to be a long-living process that precedes the
+//! launch of the Exo application, and responsible for ensuring the system (configuration, settings,
+//! etc.) is in an appropriate state to facilitate the running of Exo application.
+//! The **_System Custodian_** daemon shall expose a [D-Bus](https://www.freedesktop.org/wiki/Software/dbus/)
+//! service which Exo application use to _control & query_ it.
+//!
+//! # Lifecycle
+//! When the Exo application starts, it will _wake_ the **_System Custodian_** daemon for the
+//! duration of its lifetime, and after it has terminated the daemon will go back to sleep. When
+//! the daemon wakes up, it will configure the system into a state suitable for the Exo Application;
+//! When the daemon goes to sleep, it will revert those changes as much as it can in case they were
+//! destructive to the user's pre-existing configurations.
+//!
+//! # Responsibilities
+//! TODO: these are purely on MacOS, but change to be more broad
+//! The **_System Custodian_** daemon is responsible for using System Configuration framework to
+//!  1. duplicate the current network set
+//!  2. modify existing services to turn on IPv6 if not there
+//!  3. remove any bridge services & add any missing services that AREN'T bridge
+//! TODO: In the future:
+//!  1. run a dummy AWDL service to [allow for macOS peer-to-peer wireless networking](https://yggdrasil-network.github.io/2019/08/19/awdl.html)
+//!  2. toggle some GPU/memory configurations to speed up GPU (ask Alex what those configurations are)
+//!  3. if we ever decide to provide our **own network interfaces** that abstract over some userland
+//!     logic, this would be the place to spin that up.
+//!
+//! Then it will watch the SCDynamicStore for:
+//!  1. all __actual__ network interfaces -> collect information on them e.g. their BSD name, MAC
+//!     address, MTU, IPv6 addresses, etc. -> and set up watchers/notifiers to inform the DBus
+//!     interface of any changes
+//!  2. watch for any __undesirable__ changes to configuration and revert it
+//!
+//! It should somehow (probably through system sockets and/or BSD interface) trigger IPv6 NDP on
+//! each of the interfaces & also listen to/query for any changes on the OS routing cache??
+//! Basically emulate the `ping6 ff02::1%enX` and `ndp -an` commands BUT BETTER!!!
+//!  1. all that info should coalesce back to the overall state colleted -> should be queryable
+//!     over D-Bus
+//! TODO:
+//!  1. we might potentially add to this step a handshake of some kind...? To ensure that we can
+//!     ACTUALLY communicate with that machine over that link over e.g. TCP, UDP, etc. Will the
+//!     handshake require to know Node ID? Will the handshake require heartbeats? Who knows...
+//!  2. if we ever decide to write proprietary L2/L3 protocols for quicker communication,
+//!     e.g. [AF_NDRV](https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/)
+//!     for raw ethernet frame communication, or even a [custom thunderbolt PCIe driver](https://developer.apple.com/documentation/pcidriverkit/creating-custom-pcie-drivers-for-thunderbolt-devices),
+//!     then this would be the place to carry out discovery and propper handshakes with devices
+//!     on the other end of the link.
+//!
+
+// enable Rust-unstable features for convenience
+#![feature(trait_alias)]
+#![feature(stmt_expr_attributes)]
+#![feature(type_alias_impl_trait)]
+#![feature(specialization)]
+#![feature(unboxed_closures)]
+#![feature(const_trait_impl)]
+#![feature(fn_traits)]
+
+pub(crate) mod private {
+    // sealed traits support
+    pub trait Sealed {}
+    impl<T: ?Sized> Sealed for T {}
+}
+
+/// Namespace for all the type/trait aliases used by this crate.
+pub(crate) mod alias {}
+
+/// Namespace for crate-wide extension traits/methods
+pub(crate) mod ext {}

src/exo/main.py

@@ -1,7 +1,6 @@
 import argparse
 import multiprocessing as mp
 import os
-import resource
 import signal
 from dataclasses import dataclass, field
 from typing import Self
@@ -196,8 +195,6 @@ class Node:
 
 def main():
     args = Args.parse()
-    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-    resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
 
     mp.set_start_method("spawn")
     # TODO: Refactor the current verbosity system

src/exo/master/api.py

@@ -236,6 +236,7 @@ class API:
                     instance_meta=instance_meta,
                     min_nodes=min_nodes,
                 ),
+                node_profiles=self.state.node_profiles,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
             )
@@ -291,6 +292,7 @@ class API:
                             instance_meta=instance_meta,
                             min_nodes=min_nodes,
                         ),
+                        node_profiles=self.state.node_profiles,
                         topology=self.state.topology,
                         current_instances=self.state.instances,
                     )
@@ -599,9 +601,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 profile in self.state.node_profiles.values():
+            total_available += profile.memory.ram_available
 
         return total_available
 

src/exo/master/main.py

@@ -158,6 +158,7 @@ class Master:
                                 command,
                                 self.state.topology,
                                 self.state.instances,
+                                self.state.node_profiles,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement
@@ -200,9 +201,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([x for x in 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:

src/exo/master/placement.py

@@ -6,9 +6,10 @@ from typing import Sequence
 from loguru import logger
 
 from exo.master.placement_utils import (
+    NodeWithProfile,
     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,
@@ -19,10 +20,11 @@ from exo.shared.types.commands import (
     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 NodePerformanceProfile
 from exo.shared.types.worker.instances import (
     Instance,
     InstanceId,
@@ -52,19 +54,16 @@ def place_instance(
     command: PlaceInstance,
     topology: Topology,
     current_instances: Mapping[InstanceId, Instance],
+    node_profiles: Mapping[NodeId, NodePerformanceProfile],
 ) -> 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)
-    )
+    cycles = topology.get_cycles() + [[node] for node in all_nodes]
+    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_profiles, command.model_meta.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:
@@ -94,13 +93,15 @@ def place_instance(
     smallest_tb_cycles = [
         cycle
         for cycle in smallest_cycles
-        if topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
+        if topology.get_subgraph_from_nodes(
+            [node.node_id for node in cycle]
+        ).is_thunderbolt_cycle([node.node_id for node in cycle])
     ]
 
     if smallest_tb_cycles != []:
         smallest_cycles = smallest_tb_cycles
 
-    cycles_with_leaf_nodes: list[list[NodeInfo]] = [
+    cycles_with_leaf_nodes: list[list[NodeWithProfile]] = [
         cycle
         for cycle in smallest_cycles
         if any(topology.node_is_leaf(node.node_id) for node in cycle)
@@ -109,11 +110,7 @@ def place_instance(
     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.node_profile.memory.ram_available for node in cycle),
             start=Memory(),
         ),
     )
@@ -122,14 +119,16 @@ def place_instance(
         command.model_meta, selected_cycle, command.sharding
     )
 
-    cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle)
+    cycle_digraph: Topology = topology.get_subgraph_from_nodes(
+        [node.node_id for node in selected_cycle]
+    )
 
     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 +136,18 @@ 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(
                 cycle_digraph,
             )
             mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
-                selected_cycle,
+                coordinator=selected_cycle[0].node_id,
                 coordinator_port=random_ephemeral_port(),
                 cycle_digraph=cycle_digraph,
             )
             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:

src/exo/master/placement_utils.py

@@ -1,5 +1,4 @@
-from collections.abc import Generator
-from typing import TypeGuard, cast
+from collections.abc import Generator, Mapping
 
 from loguru import logger
 from pydantic import BaseModel
@@ -9,7 +8,7 @@ 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.topology import RDMAConnection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
@@ -24,27 +23,32 @@ class NodeWithProfile(BaseModel):
     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[list[NodeId]],
+    node_profiles: Mapping[NodeId, NodePerformanceProfile],
+    required_memory: Memory,
+) -> list[list[NodeWithProfile]]:
+    filtered_cycles: list[list[NodeWithProfile]] = []
     for cycle in cycles:
-        if not narrow_all_nodes(cycle):
+        if not all(node in node_profiles for node in cycle):
             continue
 
         total_mem = sum(
-            (node.node_profile.memory.ram_available for node in cycle), start=Memory()
+            (node_profiles[node].memory.ram_available for node in cycle), start=Memory()
         )
         if total_mem >= required_memory:
-            filtered_cycles.append(cast(list[NodeInfo], cycle))
+            filtered_cycles.append(
+                [
+                    NodeWithProfile(node_id=node, node_profile=node_profiles[node])
+                    for node in cycle
+                ]
+            )
     return filtered_cycles
 
 
-def get_smallest_cycles(cycles: list[list[NodeInfo]]) -> list[list[NodeInfo]]:
+def get_smallest_cycles(
+    cycles: list[list[NodeWithProfile]],
+) -> list[list[NodeWithProfile]]:
     min_nodes = min(len(cycle) for cycle in cycles)
     return [cycle for cycle in cycles if len(cycle) == min_nodes]
 
@@ -135,11 +139,9 @@ def get_shard_assignments_for_tensor_parallel(
 
 def get_shard_assignments(
     model_meta: ModelMetadata,
-    selected_cycle: list[NodeInfo],
+    selected_cycle: list[NodeWithProfile],
     sharding: Sharding,
 ) -> 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(
@@ -176,17 +178,16 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
         current_node = cycle[i]
         next_node = cycle[(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
-            ):
+        for src, sink, connection in cycle_digraph.list_connections():
+            if not isinstance(connection, SocketConnection):
+                continue
+
+            if src == current_node and sink == next_node:
                 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,
+                    ip=connection.sink_multiaddr.ip_address,
+                    port=connection.sink_multiaddr.port,
                 )
                 hosts.append(host)
                 break
@@ -194,8 +195,7 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
     return hosts
 
 
-def get_mlx_ibv_devices_matrix(
-    selected_cycle: list[NodeInfo],
+def get_mlx_jaccl_devices_matrix(
     cycle_digraph: Topology,
 ) -> list[list[str | None]]:
     """Build connectivity matrix mapping device i to device j via RDMA interface names.
@@ -204,6 +204,7 @@ def get_mlx_ibv_devices_matrix(
     to device j, or None if no connection exists or no interface name is found.
     Diagonal elements are always None.
     """
+    selected_cycle = list(cycle_digraph.list_nodes())
     num_nodes = len(selected_cycle)
     matrix: list[list[str | None]] = [
         [None for _ in range(num_nodes)] for _ in range(num_nodes)
@@ -214,71 +215,38 @@ 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."""
+    # TODO: Prioritise ETHERNET > ??WIFI > TB for coordinator
+    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,
+    node_info: NodeWithProfile,
 ) -> 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:
         if interface.ip_address == ip_address:
             return interface.name
@@ -287,7 +255,7 @@ def _find_interface_name_for_ip(
 
 
 def _find_ip_prioritised(
-    node: NodeInfo, other_node: NodeInfo, cycle_digraph: Topology
+    node: NodeWithProfile, other_node: NodeWithProfile, cycle_digraph: Topology
 ) -> str | None:
     # TODO: Actually prioritize in the correct Ethernet > Wifi > Non-TB > TB order.
     """Find an IP address between nodes with prioritization.
@@ -298,7 +266,7 @@ 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.node_id, other_node.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}
 
@@ -324,7 +292,7 @@ def _find_ip_prioritised(
 
 
 def get_mlx_ring_hosts_by_node(
-    selected_cycle: list[NodeInfo],
+    selected_cycle: list[NodeWithProfile],
     cycle_digraph: Topology,
     ephemeral_port: int,
 ) -> dict[NodeId, list[Host]]:
@@ -361,7 +329,7 @@ def get_mlx_ring_hosts_by_node(
             connection_ip = _find_ip_prioritised(node, other_node, cycle_digraph)
             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}"
                 )
                 raise ValueError(
                     "MLX ring backend requires connectivity between neighbouring nodes"
@@ -375,31 +343,30 @@ def get_mlx_ring_hosts_by_node(
 
 
 def get_mlx_jaccl_coordinators(
-    selected_cycle: list[NodeInfo],
+    coordinator: NodeId,
     coordinator_port: int,
     cycle_digraph: Topology,
 ) -> 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}")
+    selected_cycle = list(cycle_digraph.list_nodes())
+    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:
+        for ip, _ in _find_connection_ip(n, coordinator, cycle_digraph):
             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
-    }
+    return {n: f"{get_ip_for_node(n)}:{coordinator_port}" for n in selected_cycle}

src/exo/master/tests/conftest.py

@@ -1,67 +1,36 @@
-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,
+    MemoryUsage,
     NodePerformanceProfile,
     SystemPerformanceProfile,
 )
-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_profile(memory: int) -> NodePerformanceProfile:
+    return NodePerformanceProfile(
+        model_id="test",
+        chip_id="test",
+        friendly_name="test",
+        memory=MemoryUsage.from_bytes(
+            ram_total=1000,
+            ram_available=memory,
+            swap_total=1000,
+            swap_available=1000,
+        ),
+        network_interfaces=[],
+        system=SystemPerformanceProfile(),
+    )
 
 
 # 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_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_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_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_master.py

@@ -19,15 +19,13 @@ 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
@@ -83,21 +81,14 @@ async def test_master():
                 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),
                         ),
                     )
                 ),
@@ -163,7 +154,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)

src/exo/master/tests/test_placement.py

@@ -1,5 +1,3 @@
-from typing import Callable
-
 import pytest
 from loguru import logger
 
@@ -7,14 +5,20 @@ from exo.master.placement import (
     get_transition_events,
     place_instance,
 )
+from exo.master.tests.conftest import (
+    create_connection,
+    create_node_profile,
+    create_rdma_connection,
+)
 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
+from exo.shared.types.topology import 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(
@@ -77,34 +76,36 @@ 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],
 ):
     # arrange
     model_meta.n_layers = total_layers
     model_meta.storage_size.in_bytes = sum(
         available_memory
     )  # make it exactly fit across all nodes
+    topology = Topology()
 
     cic = place_instance_command(model_meta)
     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))
+    profiles = {
+        node_id_a: create_node_profile(available_memory[0]),
+        node_id_b: create_node_profile(available_memory[1]),
+        node_id_c: create_node_profile(available_memory[2]),
+    }
+    topology.add_node(node_id_a)
+    topology.add_node(node_id_b)
+    topology.add_node(node_id_c)
+    topology.add_connection(node_id_a, node_id_b, create_connection(1))
+    topology.add_connection(node_id_b, node_id_c, create_connection(2))
+    topology.add_connection(node_id_c, node_id_a, create_connection(3))
+    topology.add_connection(node_id_c, node_id_b, create_connection(4))
+    topology.add_connection(node_id_a, node_id_c, create_connection(5))
+    topology.add_connection(node_id_b, node_id_a, create_connection(6))
 
     # act
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, profiles)
 
     # assert
     assert len(placements) == 1
@@ -130,12 +131,11 @@ 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)
+    profiles = {node_id: create_node_profile(1000 * 1024)}
     cic = place_instance_command(
         ModelMetadata(
             model_id=ModelId("test-model"),
@@ -146,7 +146,7 @@ def test_get_instance_placements_one_node_exact_fit(
             supports_tensor=True,
         ),
     )
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, profiles)
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -157,12 +157,11 @@ 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)
+    profiles = {node_id: create_node_profile(1001 * 1024)}
     cic = place_instance_command(
         ModelMetadata(
             model_id=ModelId("test-model"),
@@ -173,7 +172,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory(
             supports_tensor=True,
         ),
     )
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, profiles)
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -184,12 +183,11 @@ 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)
+    profiles = {node_id: create_node_profile(1000 * 1024)}
     cic = place_instance_command(
         model_meta=ModelMetadata(
             model_id=ModelId("test-model"),
@@ -202,7 +200,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, {}, profiles)
 
 
 def test_get_transition_events_no_change(instance: Instance):
@@ -247,179 +245,103 @@ 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,
+def test_placement_selects_leaf_nodes(
     model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
 ):
-    # 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_meta.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))
+    profiles = {
+        node_id_a: create_node_profile(500),
+        node_id_b: create_node_profile(600),
+        node_id_c: create_node_profile(600),
+        node_id_d: create_node_profile(500),
+    }
 
-    # 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))
+    # Daisy chain topology
+    topology.add_connection(node_id_a, node_id_b, create_connection(1))
+    topology.add_connection(node_id_b, node_id_a, create_connection(1))
+    topology.add_connection(node_id_b, node_id_c, create_connection(1))
+    topology.add_connection(node_id_c, node_id_b, create_connection(1))
+    topology.add_connection(node_id_c, node_id_d, create_connection(1))
+    topology.add_connection(node_id_d, node_id_c, create_connection(1))
 
-    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))
+    logger.info(list(topology.list_connections()))
 
     cic = place_instance_command(
         model_meta=model_meta,
     )
 
-    # Act
-    placements = place_instance(cic, topology, {})
+    # act
+    placements = place_instance(cic, topology, {}, profiles)
 
-    # Assert: D-E-F cycle should be selected as it has more total memory
+    # 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],
 ):
+    topology = Topology()
     model_meta.n_layers = 12
     model_meta.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)
+    profiles = {
+        node_a: create_node_profile(500),
+        node_b: create_node_profile(500),
+        node_c: create_node_profile(500),
+    }
 
     ethernet_interface = NetworkInterfaceInfo(
         name="en0",
         ip_address="192.168.1.100",
     )
-
-    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,
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address=f"/ip4/192.168.1.{100}/tcp/{8000}")
     )
 
+    profiles[node_a].network_interfaces = [ethernet_interface]
+    profiles[node_b].network_interfaces = [ethernet_interface]
+    profiles[node_c].network_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)
+    topology.add_connection(node_a, node_b, create_rdma_connection(3))
+    topology.add_connection(node_b, node_c, create_rdma_connection(4))
+    topology.add_connection(node_c, node_a, create_rdma_connection(5))
+    topology.add_connection(node_b, node_a, create_rdma_connection(3))
+    topology.add_connection(node_c, node_b, create_rdma_connection(4))
+    topology.add_connection(node_a, node_c, create_rdma_connection(5))
+
+    topology.add_connection(node_a, node_b, ethernet_conn)
+    topology.add_connection(node_b, node_c, ethernet_conn)
+    topology.add_connection(node_c, node_a, ethernet_conn)
+    topology.add_connection(node_a, node_c, ethernet_conn)
+    topology.add_connection(node_b, node_a, ethernet_conn)
+    topology.add_connection(node_c, node_b, ethernet_conn)
 
     cic = PlaceInstance(
         sharding=Sharding.Tensor,
@@ -429,7 +351,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
         min_nodes=1,
     )
 
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, profiles)
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -437,10 +359,10 @@ def test_tensor_rdma_backend_connectivity_matrix(
 
     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):
@@ -449,15 +371,15 @@ def test_tensor_rdma_backend_connectivity_matrix(
     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

src/exo/master/tests/test_placement_utils.py

@@ -1,56 +1,48 @@
-from typing import Callable
-
 import pytest
 
 from exo.master.placement_utils import (
+    NodeWithProfile,
     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_connection, create_node_profile
 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.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()
+    topology = Topology()
 
-    node1 = create_node(1000 * 1024, node1_id)
-    node2 = create_node(1000 * 1024, node2_id)
+    node1 = create_node_profile(1000 * 1024)
+    node2 = create_node_profile(1000 * 1024)
+    node_profiles = {node1_id: node1, node2_id: node2}
 
-    topology.add_node(node1)
-    topology.add_node(node2)
+    topology.add_node(node1_id)
+    topology.add_node(node2_id)
 
-    connection1 = create_connection(node1_id, node2_id)
-    connection2 = create_connection(node2_id, node1_id)
+    connection1 = create_connection(1)
+    connection2 = create_connection(2)
 
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
+    topology.add_connection(node1_id, node2_id, connection1)
+    topology.add_connection(node2_id, node1_id, connection2)
 
     cycles = topology.get_cycles()
     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_profiles, Memory.from_bytes(1)
+    )
 
     # assert
     assert len(filtered_cycles) == 1
@@ -58,64 +50,65 @@ def test_filter_cycles_by_memory(
     assert set(n.node_id 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()
+    topology = Topology()
 
-    node1 = create_node(1000 * 1024, node1_id)
-    node2 = create_node(1000 * 1024, node2_id)
+    node1 = create_node_profile(1000 * 1024)
+    node2 = create_node_profile(1000 * 1024)
+    node_profiles = {node1_id: node1, node2_id: node2}
 
-    topology.add_node(node1)
-    topology.add_node(node2)
+    topology.add_node(node1_id)
+    topology.add_node(node2_id)
 
-    connection1 = create_connection(node1_id, node2_id)
-    connection2 = create_connection(node2_id, node1_id)
+    connection1 = create_connection(1)
+    connection2 = create_connection(2)
 
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
+    topology.add_connection(node1_id, node2_id, connection1)
+    topology.add_connection(node2_id, node1_id, connection2)
 
     # act
     filtered_cycles = filter_cycles_by_memory(
-        topology.get_cycles(), Memory.from_kb(2001)
+        topology.get_cycles(), node_profiles, 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()
+    topology = Topology()
 
-    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 = create_node_profile(500 * 1024)
+    node_b = create_node_profile(500 * 1024)
+    node_c = create_node_profile(1000 * 1024)
+    node_profiles = {
+        node_a_id: node_a,
+        node_b_id: node_b,
+        node_c_id: node_c,
+    }
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
 
-    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.add_connection(node_a_id, node_b_id, create_connection(1))
+    topology.add_connection(node_b_id, node_a_id, create_connection(2))
+    topology.add_connection(node_a_id, node_c_id, create_connection(3))
+    topology.add_connection(node_c_id, node_b_id, create_connection(4))
 
     cycles = topology.get_cycles()
 
     # act
-    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_kb(1500))
+    filtered_cycles = filter_cycles_by_memory(
+        cycles, node_profiles, Memory.from_kb(1500)
+    )
 
     # assert
     assert len(filtered_cycles) == 1
@@ -127,31 +120,38 @@ def test_filter_multiple_cycles_by_memory(
     }
 
 
-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()
+    topology = Topology()
 
-    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 = create_node_profile(500 * 1024)
+    node_b = create_node_profile(500 * 1024)
+    node_c = create_node_profile(1000 * 1024)
+    node_profiles = {
+        node_a_id: node_a,
+        node_b_id: node_b,
+        node_c_id: node_c,
+    }
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
 
-    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(node_a_id, node_b_id, create_connection(1))
+    topology.add_connection(node_b_id, node_a_id, create_connection(2))
+    topology.add_connection(node_a_id, node_c_id, create_connection(3))
+    topology.add_connection(node_c_id, node_b_id, create_connection(4))
+
+    cycles = [
+        [NodeWithProfile(node_id=nid, node_profile=node_profiles[nid]) for nid in cycle]
+        for cycle in topology.get_cycles()
+    ]
 
     # act
-    smallest_cycles = get_smallest_cycles(topology.get_cycles())
+    smallest_cycles = get_smallest_cycles(cycles)
 
     # assert
     assert len(smallest_cycles) == 1
@@ -168,9 +168,6 @@ def test_get_smallest_cycles(
     ],
 )
 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],
@@ -179,19 +176,25 @@ def test_get_shard_assignments(
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
+    topology = Topology()
 
-    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)
+    node_a = create_node_profile(available_memory[0] * 1024)
+    node_b = create_node_profile(available_memory[1] * 1024)
+    node_c = create_node_profile(available_memory[2] * 1024)
+    node_profiles = {
+        node_a_id: node_a,
+        node_b_id: node_b,
+        node_c_id: node_c,
+    }
 
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
 
-    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(node_a_id, node_b_id, create_connection(1))
+    topology.add_connection(node_b_id, node_c_id, create_connection(2))
+    topology.add_connection(node_c_id, node_a_id, create_connection(3))
+    topology.add_connection(node_b_id, node_a_id, create_connection(4))
 
     model_meta = ModelMetadata(
         model_id=ModelId("test-model"),
@@ -201,7 +204,11 @@ def test_get_shard_assignments(
         hidden_size=1000,
         supports_tensor=True,
     )
-    cycles = topology.get_cycles()
+
+    cycles = [
+        [NodeWithProfile(node_id=nid, node_profile=node_profiles[nid]) for nid in cycle]
+        for cycle in topology.get_cycles()
+    ]
     selected_cycle = cycles[0]
 
     # act
@@ -230,28 +237,21 @@ def test_get_shard_assignments(
     )
 
 
-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)
-
-    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(node_a_id, node_b_id, create_connection(1))
+    topology.add_connection(node_b_id, node_a_id, create_connection(2))
+    topology.add_connection(node_a_id, node_c_id, create_connection(3))
+    topology.add_connection(node_c_id, node_b_id, create_connection(4))
 
     # act
     hosts = get_hosts_from_subgraph(topology)
@@ -259,108 +259,47 @@ 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.2"), port=1234),
+        Host(ip=("169.254.0.3"), port=1234),
+        Host(ip=("169.254.0.4"), 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()
+    topology = Topology()
 
-    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.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(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)
+    topology.add_connection(node_a_id, node_b_id, create_connection(1))
+    topology.add_connection(node_b_id, node_a_id, create_connection(2))
+    topology.add_connection(node_a_id, node_c_id, create_connection(3))
+    topology.add_connection(node_c_id, node_b_id, create_connection(4))
 
-    # 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
+    conn_a_b = create_connection(1)
+    conn_b_a = create_connection(2)
+    conn_b_c = create_connection(3)
+    conn_c_b = create_connection(4)
+    conn_c_a = create_connection(5)
+    conn_a_c = create_connection(6)
 
-    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,
-    )
-    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,
-    )
-    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,
-    )
-
-    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_a)
-    topology.add_connection(conn_b_c)
-    topology.add_connection(conn_c_b)
-    topology.add_connection(conn_c_a)
-    topology.add_connection(conn_a_c)
-
-    cycle = [node_a, node_b, node_c]
+    topology.add_connection(node_a_id, node_b_id, conn_a_b)
+    topology.add_connection(node_b_id, node_a_id, conn_b_a)
+    topology.add_connection(node_b_id, node_c_id, conn_b_c)
+    topology.add_connection(node_c_id, node_b_id, conn_c_b)
+    topology.add_connection(node_c_id, node_a_id, conn_c_a)
+    topology.add_connection(node_a_id, node_c_id, conn_a_c)
 
     # act
     coordinators = get_mlx_jaccl_coordinators(
-        cycle, coordinator_port=5000, cycle_digraph=topology
+        node_a_id, coordinator_port=5000, cycle_digraph=topology
     )
 
     # assert
@@ -389,11 +328,11 @@ def test_get_mlx_jaccl_coordinators(
 
     # 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"
-    ), "node_b should use the IP from conn_b_a"
+    assert coordinators[node_b_id] == (f"{conn_b_a.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 coordinators[node_c_id] == (
-        f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
-    ), "node_c should use the IP from conn_c_a"
+    assert coordinators[node_c_id] == (f"{conn_c_a.sink_multiaddr.ip_address}:5000"), (
+        "node_c should use the IP from conn_c_a"
+    )

src/exo/master/tests/test_topology.py

@@ -1,13 +1,14 @@
 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,
+    MemoryUsage,
     NodePerformanceProfile,
     SystemPerformanceProfile,
 )
-from exo.shared.types.topology import Connection, ConnectionProfile, NodeId, NodeInfo
+from exo.shared.types.topology import SocketConnection
 
 
 @pytest.fixture
@@ -16,20 +17,15 @@ 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 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(
+    memory_profile = MemoryUsage.from_bytes(
         ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
     )
     system_profile = SystemPerformanceProfile()
@@ -43,162 +39,85 @@ def node_profile() -> NodePerformanceProfile:
     )
 
 
-@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, 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()
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(node_a, node_b, connection)
 
     # act
-    data = topology.get_connection_profile(connection)
+    data = list(conn for _, _, conn in 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, 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()
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(node_a, node_b, connection)
 
     # act
-    topology.remove_connection(connection)
+    topology.remove_connection(node_a, node_b, connection)
 
     # 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
-):
+def test_remove_node_still_connected(topology: Topology, 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()
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(node_a, node_b, connection)
+    assert list(topology.out_edges(node_a)) == [(node_b, connection)]
 
     # 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, 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()
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(node_a, node_b, connection)
+    assert list(topology.out_edges(node_a)) == [(node_b, connection)]
 
     # 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 {node for node in nodes} == {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,23 @@ from exo.shared.types.events import (
     TopologyEdgeCreated,
     TopologyEdgeDeleted,
 )
-from exo.shared.types.profiling import NodePerformanceProfile, SystemPerformanceProfile
+from exo.shared.types.profiling import NodePerformanceProfile
 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 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,
+    MacTBConnections,
+    MacTBIdentifiers,
+    MemoryUsage,
+    MiscData,
+    NodeConfig,
+    NodeNetworkInterfaces,
+    StaticNodeInformation,
+)
 
 
 def event_apply(event: Event, state: State) -> State:
@@ -47,16 +55,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,7 +192,7 @@ def apply_runner_deleted(event: RunnerDeleted, state: State) -> State:
 
 
 def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
-    topology = copy.copy(state.topology)
+    topology = copy.deepcopy(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
@@ -196,8 +200,12 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     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
+    }
     return state.model_copy(
         update={
+            "downloads": downloads,
             "topology": topology,
             "node_profiles": node_profiles,
             "last_seen": last_seen,
@@ -205,103 +213,69 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     )
 
 
-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)
+def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
+    topology = copy.deepcopy(state.topology)
+    topology.add_node(event.node_id)
+    info = event.info
+    profile = state.node_profiles.get(event.node_id, NodePerformanceProfile())
+    # TODO: should be broken up into individual events instead of this monster
+    match info:
+        case MacmonMetrics():
+            profile.system = info.system_profile
+            profile.memory = info.memory
+        case MemoryUsage():
+            profile.memory = info
+        case NodeConfig():
+            pass
+        case MiscData():
+            profile.friendly_name = info.friendly_name
+        case StaticNodeInformation():
+            profile.model_id = info.model
+            profile.chip_id = info.chip
+        # TODO: makes me slightly sad
+        case NodeNetworkInterfaces():
+            profile.network_interfaces = info.ifaces
+        case MacTBIdentifiers():
+            profile.tb_interfaces = info.idents
+        case MacTBConnections():
+            conn_map = {
+                tb_ident.domain_uuid: (nid, tb_ident.rdma_interface)
+                for nid in state.node_profiles
+                for tb_ident in state.node_profiles[nid].tb_interfaces
+            }
+            as_rdma_conns = [
+                (
+                    conn_map[tb_conn.sink_uuid][0],
+                    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_tb_connections(event.node_id, as_rdma_conns)
+
+    last_seen = {**state.last_seen, event.node_id: datetime.fromisoformat(event.when)}
+    new_profiles = {**state.node_profiles, event.node_id: profile}
     return state.model_copy(
         update={
             "node_profiles": new_profiles,
-            "topology": topology,
             "last_seen": last_seen,
+            "topology": topology,
         }
     )
 
 
-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] = {
-            **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,
-    }
-    # 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}
-    )
-
-
-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})
-
-
 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.source, event.sink, event.edge)
     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.sink, event.source, event.edge)
     # 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

@@ -24,6 +24,8 @@ class _InterceptHandler(logging.Handler):
         except ValueError:
             level = record.levelno
 
+        return
+
         logger.opt(depth=3, exception=record.exc_info).log(level, record.getMessage())
 
 

src/exo/shared/models/model_cards.py

@@ -14,6 +14,32 @@ class ModelCard(CamelCaseModel):
 
 MODEL_CARDS: dict[str, ModelCard] = {
     # deepseek v3
+    # "deepseek-v3-0324:4bit": ModelCard(
+    #     short_id="deepseek-v3-0324:4bit",
+    #     model_id="mlx-community/DeepSeek-V3-0324-4bit",
+    #     name="DeepSeek V3 0324 (4-bit)",
+    #     description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-V3-0324-4bit"),
+    #         pretty_name="DeepSeek V3 0324 (4-bit)",
+    #         storage_size=Memory.from_kb(409706307),
+    #         n_layers=61,
+    #     ),
+    # ),
+    # "deepseek-v3-0324": ModelCard(
+    #     short_id="deepseek-v3-0324",
+    #     model_id="mlx-community/DeepSeek-v3-0324-8bit",
+    #     name="DeepSeek V3 0324 (8-bit)",
+    #     description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-v3-0324-8bit"),
+    #         pretty_name="DeepSeek V3 0324 (8-bit)",
+    #         storage_size=Memory.from_kb(754706307),
+    #         n_layers=61,
+    #     ),
+    # ),
     "deepseek-v3.1-4bit": ModelCard(
         short_id="deepseek-v3.1-4bit",
         model_id=ModelId("mlx-community/DeepSeek-V3.1-4bit"),
@@ -44,6 +70,63 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
+    # "deepseek-v3.2": ModelCard(
+    #     short_id="deepseek-v3.2",
+    #     model_id=ModelId("mlx-community/DeepSeek-V3.2-8bit"),
+    #     name="DeepSeek V3.2 (8-bit)",
+    #     description="""DeepSeek V3.2 is a large language model trained on the DeepSeek V3.2 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-V3.2-8bit"),
+    #         pretty_name="DeepSeek V3.2 (8-bit)",
+    #         storage_size=Memory.from_kb(754706307),
+    #         n_layers=61,
+    #         hidden_size=7168,
+    #     ),
+    # ),
+    # "deepseek-v3.2-4bit": ModelCard(
+    #     short_id="deepseek-v3.2-4bit",
+    #     model_id=ModelId("mlx-community/DeepSeek-V3.2-4bit"),
+    #     name="DeepSeek V3.2 (4-bit)",
+    #     description="""DeepSeek V3.2 is a large language model trained on the DeepSeek V3.2 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-V3.2-4bit"),
+    #         pretty_name="DeepSeek V3.2 (4-bit)",
+    #         storage_size=Memory.from_kb(754706307 // 2),  # TODO !!!!!
+    #         n_layers=61,
+    #         hidden_size=7168,
+    #     ),
+    # ),
+    # deepseek r1
+    # "deepseek-r1-0528-4bit": ModelCard(
+    #     short_id="deepseek-r1-0528-4bit",
+    #     model_id="mlx-community/DeepSeek-R1-0528-4bit",
+    #     name="DeepSeek-R1-0528 (4-bit)",
+    #     description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-R1-0528-4bit"),
+    #         pretty_name="DeepSeek R1 671B (4-bit)",
+    #         storage_size=Memory.from_kb(409706307),
+    #         n_layers=61,
+    #         hidden_size=7168,
+    #     ),
+    # ),
+    # "deepseek-r1-0528": ModelCard(
+    #     short_id="deepseek-r1-0528",
+    #     model_id="mlx-community/DeepSeek-R1-0528-8bit",
+    #     name="DeepSeek-R1-0528 (8-bit)",
+    #     description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-R1-0528-8bit"),
+    #         pretty_name="DeepSeek R1 671B (8-bit)",
+    #         storage_size=Memory.from_bytes(754998771712),
+    #         n_layers=61,
+    # .       hidden_size=7168,
+    #     ),
+    # ),
     # kimi k2
     "kimi-k2-instruct-4bit": ModelCard(
         short_id="kimi-k2-instruct-4bit",
@@ -440,9 +523,8 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
-    # glm 4.5
+    # Needs to be quantized g32 or g16.
     "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",
@@ -472,81 +554,19 @@ MODEL_CARDS: dict[str, ModelCard] = {
             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,
-        ),
-    ),
-    "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,
-        ),
-    ),
-    "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,
-        ),
-    ),
-    # 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,
-        ),
-    ),
-    "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,
-        ),
-    ),
+    # "devstral-2-123b-instruct-2512-8bit": ModelCard(
+    #     short_id="devstral-2-123b-instruct-2512-8bit",
+    #     model_id=ModelId("mlx-community/Devstral-2-123B-Instruct-2512-8bit"),
+    #     name="Devstral 2 123B Instruct 2512 (8-bit, MLX)",
+    #     description="""Mistral AI's Devstral 2 123B Instruct (2512) is an agentic coding model.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/Devstral-2-123B-Instruct-2512-8bit"),
+    #         pretty_name="Devstral 2 123B Instruct 2512 (8-bit, MLX)",
+    #         storage_size=Memory.from_kb(133_000_000),
+    #         n_layers=88,
+    #         hidden_size=12288,
+    #         supports_tensor=True,
+    #     ),
+    # ),
 }

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

@@ -2,7 +2,6 @@ from exo.shared.apply import apply_node_download_progress
 from exo.shared.tests.conftest import get_pipeline_shard_metadata
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import NodeDownloadProgress
-from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.worker.tests.constants import MODEL_A_ID, MODEL_B_ID
@@ -14,7 +13,6 @@ def test_apply_node_download_progress():
     event = DownloadCompleted(
         node_id=NodeId("node-1"),
         shard_metadata=shard1,
-        total_bytes=Memory(),
     )
 
     new_state = apply_node_download_progress(
@@ -30,12 +28,10 @@ def test_apply_two_node_download_progress():
     event1 = DownloadCompleted(
         node_id=NodeId("node-1"),
         shard_metadata=shard1,
-        total_bytes=Memory(),
     )
     event2 = DownloadCompleted(
         node_id=NodeId("node-1"),
         shard_metadata=shard2,
-        total_bytes=Memory(),
     )
     state = State(downloads={NodeId("node-1"): [event1]})
 
@@ -43,7 +39,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 SocketConnection
 
 
 def test_state_serialization_roundtrip() -> None:
@@ -11,17 +11,16 @@ def test_state_serialization_roundtrip() -> None:
     node_a = NodeId("node-a")
     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"),
+    connection = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
     )
 
     state = State()
-    state.topology.add_connection(connection)
+    state.topology.add_connection(node_a, node_b, connection)
 
     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,215 @@
 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 RDMAConnection, SocketConnection
 
 
 class TopologySnapshot(BaseModel):
-    nodes: list[NodeInfo]
-    connections: list[Connection]
+    nodes: Sequence[NodeId]
+    connections: Iterable[tuple[NodeId, NodeId, 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()
+    # the _graph can be used as a int -> NodeId map.
+    _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.list_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, sink, conn in snapshot.connections:
+            topology.add_connection(source, sink, conn)
 
         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[tuple[NodeId, SocketConnection | RDMAConnection]]:
+        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 (
+            (self._graph[nid], conn)
+            for _, nid, conn 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
-
-    def contains_connection(self, connection: Connection) -> bool:
-        return connection in self._edge_id_to_rx_id_map
+        return node_id in self._vertex_indices
 
     def add_connection(
         self,
-        connection: Connection,
+        source: NodeId,
+        sink: NodeId,
+        connection: SocketConnection | RDMAConnection,
     ) -> 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:
+        if connection 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, connection)
 
-    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[tuple[NodeId, NodeId, SocketConnection | RDMAConnection]]:
+        return (
+            (
+                self._graph[src_id],
+                self._graph[sink_id],
+                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_tb_connections(
+        self, source: NodeId, new_connections: Sequence[tuple[NodeId, RDMAConnection]]
+    ) -> 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 sink, conn in new_connections:
+            self.add_connection(source, sink, conn)
+
+    def remove_connection(
+        self, source: NodeId, sink: NodeId, edge: SocketConnection | RDMAConnection
+    ) -> None:
+        if source not in self._vertex_indices or 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[source], self._vertex_indices[sink]
+        ):
+            if self._graph.get_edge_data_by_index(conn_idx) == edge:
+                self._graph.remove_edge_from_index(conn_idx)
 
-    def get_cycles(self) -> list[list[NodeInfo]]:
+    def get_cycles(self) -> list[list[NodeId]]:
         cycle_idxs = rx.simple_cycles(self._graph)
-        cycles: list[list[NodeInfo]] = []
+        cycles: list[list[NodeId]] = []
         for cycle_idx in cycle_idxs:
             cycle = [self._graph[idx] for idx in cycle_idx]
             cycles.append(cycle)
 
         return cycles
 
-    def get_cycles_tb(self) -> list[list[NodeInfo]]:
+    def get_cycles_tb(self) -> list[list[NodeId]]:
         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[list[NodeId]] = []
         for cycle_idx in cycle_idxs:
             cycle = [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":
+        rx_idxs = [self._vertex_indices[idx] for idx in node_ids]
         topology = Topology()
         for rx_idx in rx_idxs:
             topology.add_node(self._graph[rx_idx])
-        for connection in self.list_connections():
-            if (
-                connection.local_node_id in node_idxs
-                and connection.send_back_node_id in node_idxs
-            ):
-                topology.add_connection(connection)
+        for source, sink, connection in self.list_connections():
+            if source in node_ids and sink in node_ids:
+                topology.add_connection(source, sink, 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: list[NodeId]) -> 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/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 SocketConnection
 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 naem
+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  # NB: this model is UNTAGGED!!! be warned for ser/de errors.
 
 
 class NodeDownloadProgress(BaseEvent):
@@ -107,11 +97,15 @@ class ChunkGenerated(BaseEvent):
 
 
 class TopologyEdgeCreated(BaseEvent):
-    edge: Connection
+    source: NodeId
+    sink: NodeId
+    edge: SocketConnection
 
 
 class TopologyEdgeDeleted(BaseEvent):
-    edge: Connection
+    source: NodeId
+    sink: NodeId
+    edge: SocketConnection
 
 
 Event = (
@@ -125,10 +119,8 @@ Event = (
     | InstanceDeleted
     | RunnerStatusUpdated
     | RunnerDeleted
-    | NodeCreated
     | NodeTimedOut
-    | NodePerformanceMeasured
-    | NodeMemoryMeasured
+    | NodeGatheredInfo
     | NodeDownloadProgress
     | ChunkGenerated
     | TopologyEdgeCreated

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 TBIdentifier
 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):
@@ -53,15 +54,16 @@ class NetworkInterfaceInfo(CamelCaseModel):
 
 
 class NodePerformanceProfile(CamelCaseModel):
-    model_id: str
-    chip_id: str
-    friendly_name: str
-    memory: MemoryPerformanceProfile
-    network_interfaces: list[NetworkInterfaceInfo] = []
-    system: SystemPerformanceProfile
+    model_id: str = "Unknown"
+    chip_id: str = "Unknown"
+    friendly_name: str = "Unknown"
+    memory: MemoryUsage = MemoryUsage.from_bytes(
+        ram_total=0, ram_available=0, swap_total=0, swap_available=0
+    )
+    network_interfaces: Sequence[NetworkInterfaceInfo] = []
+    tb_interfaces: Sequence[TBIdentifier] = []
+    system: SystemPerformanceProfile = SystemPerformanceProfile()
 
 
 class ConnectionProfile(CamelCaseModel):
-    throughput: float
-    latency: float
-    jitter: float
+    pass

src/exo/shared/types/thunderbolt.py

@@ -0,0 +1,75 @@
+import anyio
+from pydantic import BaseModel, Field
+
+from exo.utils.pydantic_ext import CamelCaseModel
+
+
+class TBConnection(CamelCaseModel):
+    source_uuid: str
+    sink_uuid: str
+
+
+class TBIdentifier(CamelCaseModel):
+    rdma_interface: str
+    domain_uuid: str
+
+
+## Intentionally minimal, only collecting data we care about - there's a lot more
+
+
+class TBReceptacleTag(BaseModel, extra="ignore"):
+    receptacle_id_key: str | None = None
+
+
+class TBConnectivityItem(BaseModel, extra="ignore"):
+    domain_uuid_key: str | None = None
+
+
+class TBConnectivityData(BaseModel, extra="ignore"):
+    domain_uuid_key: str | None = None
+    items: list[TBConnectivityItem] | None = Field(None, alias="_items")
+    receptacle_1_tag: TBReceptacleTag | None = None
+
+    def ident(self, ifaces: dict[str, str]) -> TBIdentifier | 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 TBIdentifier(rdma_interface=iface, domain_uuid=self.domain_uuid_key)
+
+    def conn(self) -> TBConnection | 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 TBConnection(source_uuid=self.domain_uuid_key, sink_uuid=sink_key)
+
+
+class TBConnectivity(BaseModel, extra="ignore"):
+    SPThunderboltDataType: list[TBConnectivityData] = []
+
+    @classmethod
+    async def gather(cls) -> list[TBConnectivityData] | 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 TBConnectivity.model_validate_json(proc.stdout).SPThunderboltDataType

src/exo/shared/types/topology.py

@@ -1,37 +1,32 @@
-from exo.shared.types.common import NodeId
+from enum import Enum
+
+from loguru import logger
+
 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
-
-
-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")
+        logger.warning("duh")
+        return True
+
+
+# TODO
+class LinkType(str, Enum):
+    Thunderbolt = "Thunderbolt"
+    Ethernet = "Ethernet"
+    WiFi = "WiFi"
+
+
+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")

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

@@ -28,7 +28,7 @@ class DownloadPending(BaseDownloadProgress):
 
 
 class DownloadCompleted(BaseDownloadProgress):
-    total_bytes: Memory
+    pass
 
 
 class DownloadFailed(BaseDownloadProgress):

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/utils/info_gatherer/info_gatherer.py

@@ -0,0 +1,232 @@
+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 TBConnection, TBConnectivity, TBIdentifier
+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 MacTBIdentifiers(TaggedModel):
+    idents: Sequence[TBIdentifier]
+
+
+class MacTBConnections(TaggedModel):
+    conns: Sequence[TBConnection]
+
+
+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"""
+
+    # TODO
+    @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
+    | MacTBIdentifiers
+    | MacTBConnections
+    | 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 (macmon_path := shutil.which("macmon")) is not None:
+                tg.start_soon(self._monitor_macmon, macmon_path)
+            if IS_DARWIN:
+                tg.start_soon(self._monitor_system_profiler)
+            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(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 TBConnectivity.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(MacTBIdentifiers(idents=idents))
+            old_idents = idents
+
+            conns = [it for i in data if (it := i.conn()) is not None]
+            await self.info_sender.send(MacTBConnections(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

@@ -1,13 +1,12 @@
 import http.client
-import time
+from collections.abc import Mapping
 
 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
+from exo.shared.types.profiling import NodePerformanceProfile
 
 
 async def check_reachability(
@@ -18,9 +17,8 @@ async def check_reachability(
 ) -> 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)
+    def _fetch_remote_node_id() -> NodeId | None:
+        connection = http.client.HTTPConnection(target_ip, 52415, timeout=1)
         try:
             connection.request("GET", "/node_id")
             response = connection.getresponse()
@@ -36,16 +34,7 @@ async def check_reachability(
             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}")
+        except http.client.HTTPException:
             return None
         finally:
             connection.close()
@@ -71,19 +60,20 @@ async def check_reachability(
 
 
 async def check_reachable(
-    topology: Topology, self_node_id: NodeId
+    topology: Topology,
+    profiles: Mapping[NodeId, NodePerformanceProfile],
+    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:
+        for node_id in topology.list_nodes():
+            if node_id not in profiles:
                 continue
-            for iface in node.node_profile.network_interfaces:
+            for iface in profiles[node_id].network_interfaces:
                 tg.start_soon(
                     check_reachability,
                     iface.ip_address,
-                    node.node_id,
+                    node_id,
                     self_node_id,
                     reachable,
                 )

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

@@ -0,0 +1,18 @@
+import pytest
+import sys
+from exo.shared.types.thunderbolt import (
+    TBConnectivity,
+)
+from exo.utils.info_gatherer.info_gatherer import _gather_iface_map  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.anyio
+@pytest.mark.skipif(sys.platform != "darwin", reason="TB info can only be gathered on macos")
+async def test_tb_parsing():
+    data = await TBConnectivity.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/engines/mlx/auto_parallel.py

@@ -10,24 +10,18 @@ from mlx.nn.layers.distributed import (
     shard_linear,
     sum_gradients,
 )
+from mlx_lm.models.cache import (
+    _BaseCache,  # pyright: ignore[reportPrivateUsage]
+)
 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
-from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
 
-from exo.shared.logging import logger
-from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.shared.types.worker.shards import (
+    PipelineShardMetadata,
+)
 
 
 class _LayerCallable(Protocol):
@@ -97,6 +91,8 @@ class PipelineLastLayer(CustomMlxLayer):
             x, *args, **kwargs
         ).arguments.get("cache", None)
 
+        assert cache is None or issubclass(type(cache), _BaseCache)  # type: ignore
+
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
         if self.r != self.s - 1:
@@ -104,6 +100,7 @@ class PipelineLastLayer(CustomMlxLayer):
                 output, (self.r + 1) % self.s, group=self.group
             )
             if cache is not None:
+                # This change happened upstream - check out mlx github somewhere??
                 cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
 
         output = mx.distributed.all_gather(output, group=self.group)[-output.shape[0] :]
@@ -135,6 +132,24 @@ def _get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
     return layers
 
 
+def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
+    inner_model_instance = _inner_model(model)
+    if hasattr(inner_model_instance, "layers"):
+        inner_model_instance.layers = layers
+
+        # Update DeepSeek V3 specific parameters when layers are shrunk
+        if isinstance(model, DeepseekV3Model) and hasattr(
+            inner_model_instance, "num_layers"
+        ):
+            inner_model_instance.start_idx = 0
+            inner_model_instance.end_idx = len(layers)
+            inner_model_instance.num_layers = len(layers)
+    elif hasattr(inner_model_instance, "h"):
+        inner_model_instance.h = layers
+    else:
+        raise ValueError("Model must have either a 'layers' or 'h' attribute")
+
+
 def pipeline_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
@@ -150,7 +165,8 @@ def pipeline_auto_parallel(
     """
     inner_model_instance: nn.Module = _inner_model(model)
 
-    layers = _get_layers(inner_model_instance)
+    # Handle both model.layers and model.h cases
+    layers: list[_LayerCallable] = _get_layers(inner_model_instance)
 
     start_layer, end_layer = model_shard_meta.start_layer, model_shard_meta.end_layer
     device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
@@ -164,17 +180,6 @@ def pipeline_auto_parallel(
         group=group,
     )
 
-    if isinstance(inner_model_instance, GptOssMoeModel):
-        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"
-        )
-        inner_model_instance.ga_idx = inner_model_instance.layer_types.index(  # type: ignore
-            "full_attention"
-        )
-
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -199,44 +204,18 @@ def tensor_auto_parallel(
         group=group,
     )
 
-    segments: int = 1
-
-    def _all_to_sharded(path: str, weight: mx.array):
-        if path.endswith("bias"):
-            logger.info(f"Sharding bias for {path} - all to sharded")
-            return weight.ndim - 1, segments
-        return max(weight.ndim - 2, 0), segments
-
     all_to_sharded_linear_in_place = partial(
         shard_inplace,
-        sharding=_all_to_sharded,  # type: ignore
+        sharding="all-to-sharded",
         group=group,
     )
-
-    n = group.size()
-
-    def _sharded_to_all(path: str, weight: mx.array):
-        if path.endswith("bias"):
-            logger.info(f"Sharding bias for {path} - sharded to all")
-            weight /= n
-            return None
-        return -1, segments
-
     sharded_to_all_linear_in_place = partial(
         shard_inplace,
-        sharding=_sharded_to_all,  # type: ignore
+        sharding="sharded-to-all",
         group=group,
     )
 
-    if hasattr(model, "shard"):
-        try:
-            model.shard(group)  # type: ignore
-            return model
-        except (AttributeError, TypeError, NameError):
-            pass
-
-    if isinstance(model, (LlamaModel, Ministral3Model)):
-        logger.warning("shouldn't be hit - upstream sharding exists")
+    if isinstance(model, LlamaModel):
         tensor_parallel_sharding_strategy = LlamaShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -244,8 +223,7 @@ def tensor_auto_parallel(
             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")
+    elif isinstance(model, DeepseekV3Model):
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -253,15 +231,7 @@ def tensor_auto_parallel(
             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)):
+    elif isinstance(model, Qwen3MoeModel):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -269,15 +239,6 @@ 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)}")
 
@@ -323,38 +284,13 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
         return model
 
 
-def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
-    inner_model_instance = _inner_model(model)
-    if hasattr(inner_model_instance, "layers"):
-        inner_model_instance.layers = layers
-
-        # Update DeepSeek V3 specific parameters when layers are shrunk
-        if isinstance(
-            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel)
-        ) and hasattr(inner_model_instance, "num_layers"):
-            logger.info(
-                f"Setting num_layers to {len(layers)} for model {model.model.__class__.__name__}"
-            )
-            inner_model_instance.start_idx = 0
-            inner_model_instance.end_idx = len(layers)
-            inner_model_instance.num_layers = len(layers)
-        elif isinstance(model, Qwen3MoeModel):
-            logger.info(
-                f"Setting num_hidden_layers to {len(layers)} for model {model.model.__class__.__name__}"
-            )
-            inner_model_instance.num_hidden_layers = len(layers)
-    elif hasattr(inner_model_instance, "h"):
-        inner_model_instance.h = layers
-    else:
-        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:
+            if layer.self_attn.q_lora_rank is None:  # pyright: ignore[reportUnnecessaryComparison]
+                # Unfortunately, q_lora_rank can be None despite typing hints.
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
                 )
@@ -369,7 +305,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn.num_heads //= self.N
 
             # Shard the MLP
-            if isinstance(layer.mlp, (DeepseekV3MLP, DeepseekV32MLP)):
+            if isinstance(layer.mlp, DeepseekV3MLP):
                 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)
@@ -403,35 +339,6 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
         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:
         model = cast(Qwen3MoeModel, model)
@@ -446,13 +353,11 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
 
             # Shard the MoE. Shard in place since the MoE should be responsible
             # for aggregating the results.
-            if isinstance(
-                layer.mlp, (Qwen3MoeSparseMoeBlock, MoE, Qwen3NextSparseMoeBlock)
-            ):
+            if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
                 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 = ShardedQwenMoE(layer.mlp)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
+                layer.mlp = ShardedQwenMoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group
 
             # Shard the MLP
@@ -476,50 +381,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/cache.py

@@ -1,189 +1,104 @@
-"""KV prefix cache for reusing computed prompt prefixes across requests."""
-
-from collections import OrderedDict
+# type: ignore
+# TODO: Fix this file, including types!
 from copy import deepcopy
-from typing import Any, Callable, Protocol
+from typing import Callable
 
 import mlx.core as mx
-import numpy as np
-from mlx_lm.models.cache import trim_prompt_cache
+from mlx_lm import stream_generate
+from mlx_lm.models.cache import _BaseCache, trim_prompt_cache
+from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.constants import KEEP_KV_SIZE, KV_BITS, KV_GROUP_SIZE
 from exo.worker.engines.mlx.utils_mlx import make_kv_cache
-from exo.worker.runner.bootstrap import logger
-
-# Type alias for KV cache - the actual type is _BaseCache but it's private
-KVCacheType = Any
-
-
-class TokenizerProtocol(Protocol):
-    """Protocol for tokenizers used with KVPrefixCache."""
-
-    bos_token: str | None
-
-    def encode(self, text: str, **kwargs: bool) -> list[int]: ...
 
 
 class KVPrefixCache:
-    """Cache for common prompt prefixes to avoid re-processing.
+    def __init__(self):
+        # Only one prefix cache per runner.
+        self.prompts: list[mx.array] = []  # mx array of tokens (ints)
+        self.caches: list[list[_BaseCache]] = []
 
-    Uses LRU eviction when capacity is reached. Stores tokenized prompts
-    and their corresponding KV caches for reuse.
-    """
-
-    def __init__(self, max_size: int = 10):
-        """Initialize prefix cache.
-
-        Args:
-            max_size: Maximum number of cached entries before LRU eviction.
-        """
-        self.max_size = max_size
-        # OrderedDict maintains insertion order for LRU - most recent at end
-        # Key: token bytes, Value: (tokens as mx.array, KV cache)
-        self._cache: OrderedDict[bytes, tuple[mx.array, list[KVCacheType]]] = (
-            OrderedDict()
-        )
-
-    def _token_key(self, tokens: mx.array) -> bytes:
-        """Create hashable key from token array."""
-        return np.array(tokens.tolist(), dtype=np.int32).tobytes()
-
-    def _encode_prompt(self, tokenizer: TokenizerProtocol, prompt: str) -> mx.array:
-        """Tokenize prompt string to mx.array."""
-        add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
-            tokenizer.bos_token
-        )
-        tokenized = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
-        return mx.array(tokenized)
-
-    def _find_best_prefix(self, tokens: mx.array) -> tuple[bytes | None, int]:
-        """Find cached entry with longest matching prefix.
-
-        Returns:
-            Tuple of (cache_key, prefix_length). cache_key is None if no match found.
-        """
-        best_key: bytes | None = None
-        best_length = 0
-        target_len = tokens.shape[0]
-
-        for key, (cached_tokens, _cache) in self._cache.items():
-            prefix_len = get_prefix_length(tokens, cached_tokens)
-
-            # Exact match - return immediately
-            if prefix_len == target_len and prefix_len == cached_tokens.shape[0]:
-                return key, prefix_len
-
-            # Better prefix match
-            if prefix_len > best_length:
-                best_key = key
-                best_length = prefix_len
-
-        return best_key, best_length
+    def add_kv_cache(
+        self, tokenizer: TokenizerWrapper, prompt: str, cache: list[_BaseCache]
+    ):
+        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
+        self.prompts.append(tokenized_prompt)
+        self.caches.append(deepcopy(cache))
 
     def get_kv_cache(
         self,
         model: Model,
-        tokenizer: TokenizerProtocol,
+        tokenizer: TokenizerWrapper,
         sampler: Callable[[mx.array], mx.array],
         prompt: str,
-    ) -> tuple[list[KVCacheType], int]:
-        """Get KV cache for prompt, reusing prefix if available.
+    ) -> list[_BaseCache]:
+        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
+        max_length = len(tokenized_prompt)
 
-        Args:
-            model: The model to create cache for.
-            tokenizer: Tokenizer for encoding prompt.
-            sampler: Sampler function for prefill.
-            prompt: The prompt string to process.
+        best_snapshot_index, best_snapshot_length = None, 0
 
-        Returns:
-            Tuple of (kv_cache, tokens_reused). tokens_reused indicates how many
-            tokens were reused from cache (0 if no cache hit).
-        """
-        tokens = self._encode_prompt(tokenizer, prompt)
-        target_len = int(tokens.shape[0])
+        for i, cached_prompt in enumerate(self.prompts):
+            length = _get_prefix_length(tokenized_prompt, cached_prompt)
 
-        # Find best prefix match
-        best_key, prefix_len = self._find_best_prefix(tokens)
+            if length == max_length:
+                return self.caches[i]
 
-        if best_key is not None and prefix_len > 0:
-            cached_tokens, cached_kv = self._cache[best_key]
-            cached_len = int(cached_tokens.shape[0])
+            if length > best_snapshot_length:
+                best_snapshot_index, best_snapshot_length = i, length
 
-            # Move to end (most recently used)
-            self._cache.move_to_end(best_key)
+        if best_snapshot_index is not None:
+            prompt_cache = deepcopy(self.caches[best_snapshot_index])
+            trim_prompt_cache(prompt_cache, max_length - best_snapshot_length)
+            tokenized_prompt = tokenized_prompt[best_snapshot_index:]
 
-            if prefix_len == target_len and prefix_len == cached_len:
-                # Exact match - return deepcopy directly
-                logger.debug(f"Prefix cache: exact match, reusing {prefix_len} tokens")
-                return deepcopy(cached_kv), prefix_len
-
-            # Partial match - need to trim and/or extend
-            prompt_cache = deepcopy(cached_kv)
-
-            if cached_len > prefix_len:
-                # Cached prompt is longer - trim to prefix length
-                num_to_trim = cached_len - prefix_len
-                trim_prompt_cache(prompt_cache, num_to_trim)
-                logger.debug(
-                    f"Prefix cache: trimmed {num_to_trim} tokens from cached entry"
-                )
-
-            # Note: We don't prefill remaining tokens here - stream_generate will do it
-            # when processing the full prompt with this partial cache
-            return prompt_cache, prefix_len
-
-        # No cache hit - return fresh cache (stream_generate will prefill)
-        logger.debug(
-            f"Prefix cache: miss, will prefill {target_len} tokens during generation"
-        )
-        prompt_cache = make_kv_cache(model=model)
-
-        return prompt_cache, 0
-
-    def put(
-        self, tokenizer: TokenizerProtocol, prompt: str, cache: list[KVCacheType]
-    ) -> None:
-        """Store KV cache for prompt after generation completes.
-
-        Args:
-            tokenizer: Tokenizer for encoding prompt.
-            prompt: The prompt string that was processed.
-            cache: The KV cache to store.
-        """
-        tokens = self._encode_prompt(tokenizer, prompt)
-        key = self._token_key(tokens)
-
-        # If already in cache, just move to end
-        if key in self._cache:
-            self._cache.move_to_end(key)
-            return
-
-        # Evict LRU entry if at capacity
-        if len(self._cache) >= self.max_size:
-            evicted_key, _ = self._cache.popitem(last=False)
-            logger.debug(
-                f"Prefix cache: evicted LRU entry ({len(evicted_key)} token bytes)"
+        else:
+            prompt_cache = make_kv_cache(
+                model,
+                # max_kv_size=MAX_KV_SIZE,
+                # keep=KEEP_KV_SIZE
             )
 
-        # Store deepcopy
-        self._cache[key] = (tokens, deepcopy(cache))
-        logger.debug(f"Prefix cache: stored entry with {tokens.shape[0]} tokens")
+        prefill(model, tokenizer, sampler, tokenized_prompt, prompt_cache)
 
-    def clear(self) -> None:
-        """Clear all cached entries."""
-        self._cache.clear()
+        return prompt_cache
 
-    def __len__(self) -> int:
-        """Return number of cached entries."""
-        return len(self._cache)
+    def encode_prompt(self, tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
+        add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token
+        )
+        tokenized_prompt = tokenizer.encode(
+            prompt, add_special_tokens=add_special_tokens
+        )
+        return mx.array(tokenized_prompt)
 
 
-def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
-    """Calculate length of matching prefix between two token arrays."""
-    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
+def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
+    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]), KEEP_KV_SIZE)
     if n == 0:
         return 0
 
-    equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32)
+    equal = (prompt[:n] == cached_prompt[:n]).astype(mx.int32)
     prefix_mask = mx.cumprod(equal)  # stays 1 until first mismatch, then 0 forever
     return int(mx.sum(prefix_mask).item())
+
+
+def prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt: mx.array,
+    cache: list[_BaseCache],
+) -> None:
+    for _ in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=prompt,
+        max_tokens=0,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+    ):
+        pass

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

@@ -6,6 +6,7 @@ from mlx_lm.models.cache import KVCache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+# from exo.engines.mlx.cache import KVPrefixCache
 from exo.shared.types.api import (
     BenchChatCompletionTaskParams,
     ChatCompletionMessage,
@@ -18,7 +19,6 @@ from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
@@ -119,7 +119,6 @@ def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
     task: ChatCompletionTaskParams,
-    prefix_cache: KVPrefixCache | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -136,6 +135,8 @@ def mlx_generate(
         chat_task_data=task,
     )
 
+    caches = make_kv_cache(model=model)
+
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
         # Only sample length eos tokens
@@ -147,20 +148,6 @@ def mlx_generate(
         top_p=task.top_p if task.top_p is not None else 1.0,
     )
 
-    # Get KV cache - either from prefix cache or fresh
-    tokens_reused = 0
-    if prefix_cache is not None:
-        caches, tokens_reused = prefix_cache.get_kv_cache(
-            model=model,
-            tokenizer=tokenizer,
-            sampler=sampler,
-            prompt=prompt,
-        )
-        if tokens_reused > 0:
-            logger.info(f"Prefix cache hit: reused {tokens_reused} tokens")
-    else:
-        caches = make_kv_cache(model=model)
-
     max_tokens = task.max_tokens or MAX_TOKENS
     for out in stream_generate(
         model=model,
@@ -202,9 +189,6 @@ def mlx_generate(
         )
 
         if out.finish_reason is not None:
-            # Store in prefix cache for future reuse
-            if prefix_cache is not None:
-                prefix_cache.put(tokenizer=tokenizer, prompt=prompt, cache=caches)
             break
 
         # TODO: Do we want an mx_barrier?

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

@@ -1,23 +1,10 @@
 import json
 import os
 import resource
-import sys
 import time
 from pathlib import Path
 from typing import Any, cast
 
-# Monkey-patch for transformers 5.x compatibility
-# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
-# which was moved in transformers 5.0.0rc2
-try:
-    import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
-    from transformers.convert_slow_tokenizer import bytes_to_unicode
-
-    if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
-        gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[attr-defined]
-except ImportError:
-    pass  # transformers < 5.0 or bytes_to_unicode not available
-
 from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -31,7 +18,7 @@ from exo.worker.engines.mlx.constants import (
 try:
     from mlx_lm.tokenizer_utils import load_tokenizer
 except ImportError:
-    from mlx_lm.tokenizer_utils import load as load_tokenizer
+    from mlx_lm.tokenizer_utils import load as load_tokenizer  # type: ignore
 import contextlib
 
 import mlx.core as mx
@@ -144,20 +131,22 @@ 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_IBV_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_RANK"] = str(rank)
@@ -265,70 +254,26 @@ def shard_and_load(
     return model, tokenizer
 
 
-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)
+def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata):
+    # TODO: Let's move away from this custom logic to mlx_lm.load()
+    if "kimi-k2" in shard_metadata.model_meta.model_id.lower():
+        eos_token_ids = [163586]
 
+    elif "glm" in shard_metadata.model_meta.model_id.lower():
+        eos_token_ids = [151336, 151329, 151338]
 
-def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
-    """
-    Get the EOS token IDs for a model based on its ID.
+    else:
+        eos_token_ids = None
 
-    Some models require explicit EOS token configuration that isn't in their
-    tokenizer config. This function returns the known EOS token IDs for such models.
-
-    Args:
-        model_id: The HuggingFace model ID
-
-    Returns:
-        List of EOS token IDs, or None if the model uses standard tokenizer config
-    """
-    model_id_lower = model_id.lower()
-    if "kimi-k2" in model_id_lower:
-        return [163586]
-    elif "glm" in model_id_lower:
-        return [151336, 151329, 151338]
-    return None
-
-
-def load_tokenizer_for_model_id(model_id: str, model_path: Path) -> TokenizerWrapper:
-    """
-    Load tokenizer for a model given its ID and local path.
-
-    This is the core tokenizer loading logic, handling special cases for different
-    model families (Kimi, GLM, etc.) and transformers 5.x compatibility.
-
-    Args:
-        model_id: The HuggingFace model ID (e.g., "moonshotai/Kimi-K2-Instruct")
-        model_path: Local path where the model/tokenizer files are stored
-
-    Returns:
-        TokenizerWrapper instance configured for the model
-    """
-    model_id_lower = model_id.lower()
-    eos_token_ids = get_eos_token_ids_for_model(model_id)
-
-    # Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer
-    if "kimi-k2" in model_id_lower:
-        sys.path.insert(0, str(model_path))
-        from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
-
-        hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
-
-        # Patch encode to use internal tiktoken model directly
-        # transformers 5.x has a bug in the encode->pad path for slow tokenizers
-        def _patched_encode(text: str, **_kwargs: object) -> list[int]:
-            # Pass allowed_special="all" to handle special tokens like <|im_user|>
-            return list(hf_tokenizer.model.encode(text, allowed_special="all"))  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
-
-        hf_tokenizer.encode = _patched_encode
-        return TokenizerWrapper(hf_tokenizer, eos_token_ids=eos_token_ids)
-
-    tokenizer = load_tokenizer(
-        model_path,
-        tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
-        eos_token_ids=eos_token_ids,
+    tokenizer = cast(
+        TokenizerWrapper,
+        load_tokenizer(
+            model_path,
+            tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
+            eos_token_ids=eos_token_ids,
+        ),
     )
+    assert isinstance(tokenizer, TokenizerWrapper)
 
     return tokenizer
 
@@ -358,14 +303,14 @@ def apply_chat_template(
             {k: v for k, v in message.model_dump().items() if v is not None}  # type: ignore
         )
 
-    prompt: str = tokenizer.apply_chat_template(
+    prompt: str = tokenizer.apply_chat_template(  # type: ignore
         formatted_messages,
         tokenize=False,
         add_generation_prompt=True,
         tools=chat_task_data.tools,
     )
 
-    return prompt
+    return prompt  # type: ignore
 
 
 class NullKVCache(KVCache):

src/exo/worker/main.py

@@ -16,8 +16,7 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     NodeDownloadProgress,
-    NodeMemoryMeasured,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
     TaskCreated,
     TaskStatusUpdated,
     TopologyEdgeCreated,
@@ -25,7 +24,6 @@ from exo.shared.types.events import (
 )
 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 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
@@ -217,9 +201,7 @@ class Worker:
                     )
                     if initial_progress.status == "complete":
                         progress = DownloadCompleted(
-                            shard_metadata=shard,
-                            node_id=self.node_id,
-                            total_bytes=initial_progress.total_bytes,
+                            shard_metadata=shard, node_id=self.node_id
                         )
                         self.download_status[shard.model_meta.model_id] = progress
                         await self.event_sender.send(
@@ -252,8 +234,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 +251,24 @@ 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(
+                    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(
+                    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 +317,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
 
@@ -366,11 +346,7 @@ class Worker:
             nonlocal self
             nonlocal last_progress_time
             if progress.status == "complete":
-                status = DownloadCompleted(
-                    shard_metadata=shard,
-                    node_id=self.node_id,
-                    total_bytes=progress.total_bytes,
-                )
+                status = DownloadCompleted(shard_metadata=shard, node_id=self.node_id)
                 self.download_status[shard.model_meta.model_id] = status
                 # Footgun!
                 self.event_sender.send_nowait(
@@ -399,7 +375,6 @@ class Worker:
                 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:
@@ -422,7 +397,9 @@ class Worker:
         while True:
             # TODO: EdgeDeleted
             edges = set(self.state.topology.list_connections())
-            conns = await check_reachable(self.state.topology, self.node_id)
+            conns = await check_reachable(
+                self.state.topology, self.state.node_profiles, self.node_id
+            )
             for nid in conns:
                 for ip in conns[nid]:
                     if "127.0.0.1" in ip or "localhost" in ip:
@@ -430,26 +407,31 @@ class Worker:
                             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(
+                                source=self.node_id, sink=nid, edge=edge
+                            )
+                        )
 
             for nid, conn in self.state.topology.out_edges(self.node_id):
-                if (
-                    nid not in conns
-                    or conn.send_back_multiaddr.ip_address not in conns.get(nid, set())
+                if not isinstance(conn, SocketConnection):
+                    continue
+                if nid not in conns or conn.sink_multiaddr.ip_address not in conns.get(
+                    nid, set()
                 ):
                     logger.debug(f"ping failed to discover {conn=}")
-                    await self.event_sender.send(TopologyEdgeDeleted(edge=conn))
+                    await self.event_sender.send(
+                        TopologyEdgeDeleted(source=self.node_id, sink=nid, edge=conn)
+                    )
 
             await anyio.sleep(10)
 
@@ -463,9 +445,7 @@ class Worker:
                 ) in self.shard_downloader.get_shard_download_status():
                     if progress.status == "complete":
                         status = DownloadCompleted(
-                            node_id=self.node_id,
-                            shard_metadata=progress.shard,
-                            total_bytes=progress.total_bytes,
+                            node_id=self.node_id, shard_metadata=progress.shard
                         )
                     elif progress.status in ["in_progress", "not_started"]:
                         if progress.downloaded_bytes_this_session.in_bytes == 0:

src/exo/worker/runner/bootstrap.py

@@ -19,7 +19,7 @@ def entrypoint(
 ) -> None:
     if (
         isinstance(bound_instance.instance, MlxJacclInstance)
-        and len(bound_instance.instance.ibv_devices) >= 2
+        and len(bound_instance.instance.jaccl_devices) >= 2
     ):
         os.environ["MLX_METAL_FAST_SYNCH"] = "1"
 

src/exo/worker/runner/runner.py

@@ -39,7 +39,6 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
@@ -70,7 +69,6 @@ def main(
     model = None
     tokenizer = None
     group = None
-    prefix_cache: KVPrefixCache | None = None
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -112,8 +110,6 @@ def main(
                     )
 
                     model, tokenizer = load_mlx_items(bound_instance, group)
-                    prefix_cache = KVPrefixCache(max_size=10)
-                    logger.info("prefix cache initialized")
 
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
@@ -161,7 +157,6 @@ def main(
                         model=model,
                         tokenizer=tokenizer,
                         task=task_params,
-                        prefix_cache=prefix_cache,
                     ):
                         match response:
                             case GenerationResponse():

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

@@ -1,141 +0,0 @@
-"""Tests for KVPrefixCache."""
-
-import mlx.core as mx
-import pytest
-
-from exo.worker.engines.mlx.cache import (
-    KVCacheType,
-    KVPrefixCache,
-    TokenizerProtocol,
-    get_prefix_length,
-)
-
-
-class MockTokenizer(TokenizerProtocol):
-    """Mock tokenizer that converts string to list of char codes."""
-
-    bos_token: str | None = None
-
-    def encode(self, text: str, **kwargs: bool) -> list[int]:
-        """Encode text to list of character codes."""
-        del kwargs  # unused
-        return [ord(c) for c in text]
-
-
-class TestGetPrefixLength:
-    """Tests for the core prefix matching algorithm."""
-
-    def test_identical_arrays(self) -> None:
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3, 4, 5])
-        assert get_prefix_length(a, b) == 5
-
-    def test_partial_match(self) -> None:
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3, 9, 9])
-        assert get_prefix_length(a, b) == 3
-
-    def test_no_match(self) -> None:
-        a = mx.array([1, 2, 3])
-        b = mx.array([9, 9, 9])
-        assert get_prefix_length(a, b) == 0
-
-    def test_different_lengths(self) -> None:
-        short = mx.array([1, 2, 3])
-        long = mx.array([1, 2, 3, 4, 5])
-        # Should return length of shorter when they match
-        assert get_prefix_length(short, long) == 3
-        assert get_prefix_length(long, short) == 3
-
-    def test_empty_array(self) -> None:
-        empty: mx.array = mx.array([])
-        tokens = mx.array([1, 2, 3])
-        assert get_prefix_length(empty, tokens) == 0
-        assert get_prefix_length(tokens, empty) == 0
-
-
-class TestKVPrefixCache:
-    """Tests for the KV prefix cache."""
-
-    @pytest.fixture
-    def tokenizer(self) -> MockTokenizer:
-        """Mock tokenizer that converts string to list of char codes."""
-        return MockTokenizer()
-
-    @pytest.fixture
-    def fake_kv(self) -> list[KVCacheType]:
-        """Fake KV cache for testing."""
-        return [object()]
-
-    def test_put_stores_entry(
-        self, tokenizer: MockTokenizer, fake_kv: list[KVCacheType]
-    ) -> None:
-        cache = KVPrefixCache(max_size=10)
-
-        cache.put(tokenizer, "hello", fake_kv)
-
-        assert len(cache) == 1
-
-    def test_put_same_prompt_twice_does_not_duplicate(
-        self, tokenizer: MockTokenizer, fake_kv: list[KVCacheType]
-    ) -> None:
-        cache = KVPrefixCache(max_size=10)
-
-        cache.put(tokenizer, "hello", fake_kv)
-        cache.put(tokenizer, "hello", fake_kv)
-
-        assert len(cache) == 1
-
-    def test_lru_eviction(
-        self, tokenizer: MockTokenizer, fake_kv: list[KVCacheType]
-    ) -> None:
-        cache = KVPrefixCache(max_size=2)
-
-        # Fill cache
-        cache.put(tokenizer, "first", fake_kv)
-        cache.put(tokenizer, "second", fake_kv)
-        assert len(cache) == 2
-
-        # Add third - should evict "first" (oldest)
-        cache.put(tokenizer, "third", fake_kv)
-        assert len(cache) == 2
-
-        # Add "first" again - if it was evicted, cache size stays 2
-        # If it wasn't evicted, this would be a no-op
-        cache.put(tokenizer, "first", fake_kv)
-        # Now add fourth - if "first" was re-added, size is still 2
-        cache.put(tokenizer, "fourth", fake_kv)
-        assert len(cache) == 2
-
-    def test_lru_access_refreshes_entry(
-        self, tokenizer: MockTokenizer, fake_kv: list[KVCacheType]
-    ) -> None:
-        cache = KVPrefixCache(max_size=2)
-
-        # Add two entries
-        cache.put(tokenizer, "first", fake_kv)
-        cache.put(tokenizer, "second", fake_kv)
-
-        # Access "first" again (moves to end of LRU)
-        cache.put(tokenizer, "first", fake_kv)
-
-        # Add third - should evict "second" now (oldest)
-        cache.put(tokenizer, "third", fake_kv)
-
-        # Add "second" again - this will add it as new entry
-        cache.put(tokenizer, "second", fake_kv)
-        # Now "first" is oldest, adding fourth should evict it
-        cache.put(tokenizer, "fourth", fake_kv)
-
-        # Cache should have "second" and "fourth", not "first"
-        assert len(cache) == 2
-
-    def test_clear(self, tokenizer: MockTokenizer, fake_kv: list[KVCacheType]) -> None:
-        cache = KVPrefixCache()
-        cache.put(tokenizer, "hello", fake_kv)
-        cache.put(tokenizer, "world", fake_kv)
-        assert len(cache) == 2
-
-        cache.clear()
-
-        assert len(cache) == 0

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

@@ -1,386 +0,0 @@
-"""
-Unit tests for tokenizer loading and functionality across all supported models.
-
-This test downloads only tokenizer-related files (not full model weights) to verify
-that tokenizers can be loaded and used correctly for encoding/decoding.
-"""
-
-import asyncio
-import contextlib
-from pathlib import Path
-
-import pytest
-
-from exo.shared.models.model_cards import MODEL_CARDS, ModelCard
-from exo.worker.download.download_utils import (
-    download_file_with_retry,
-    ensure_models_dir,
-    fetch_file_list_with_cache,
-)
-from exo.worker.engines.mlx.utils_mlx import (
-    get_eos_token_ids_for_model,
-    load_tokenizer_for_model_id,
-)
-
-# Files needed for tokenizer functionality
-TOKENIZER_FILE_PATTERNS = [
-    "tokenizer.json",
-    "tokenizer_config.json",
-    "special_tokens_map.json",
-    "vocab.json",
-    "vocab.txt",
-    "merges.txt",
-    "tiktoken.model",
-    "added_tokens.json",
-    "tokenizer.model",
-    "tokenization_*.py",  # Custom tokenizer implementations
-]
-
-
-def is_tokenizer_file(filename: str) -> bool:
-    """Check if a file is needed for tokenizer functionality."""
-    for pattern in TOKENIZER_FILE_PATTERNS:
-        if "*" in pattern:
-            prefix = pattern.split("*")[0]
-            suffix = pattern.split("*")[1]
-            if filename.startswith(prefix) and filename.endswith(suffix):
-                return True
-        elif filename == pattern:
-            return True
-    return False
-
-
-async def download_tokenizer_files(model_id: str) -> Path:
-    """Download only the tokenizer-related files for a model."""
-    target_dir = await ensure_models_dir() / model_id.replace("/", "--")
-    target_dir.mkdir(parents=True, exist_ok=True)
-
-    file_list = await fetch_file_list_with_cache(model_id, "main", recursive=True)
-
-    tokenizer_files = [f for f in file_list if is_tokenizer_file(f.path)]
-
-    if not tokenizer_files:
-        pytest.skip(f"No tokenizer files found for {model_id}")
-
-    for file_entry in tokenizer_files:
-        with contextlib.suppress(FileNotFoundError):
-            await download_file_with_retry(
-                model_id, "main", file_entry.path, target_dir
-            )
-
-    return target_dir
-
-
-# Get a sample of models to test (one per family to keep tests fast)
-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():
-        # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
-        parts = short_id.split("-")
-        family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
-
-        if family not in families:
-            families[family] = (short_id, card)
-
-    return list(families.values())
-
-
-TEST_MODELS: list[tuple[str, ModelCard]] = get_test_models()
-
-
-@pytest.fixture(scope="module")
-def event_loop():
-    """Create event loop for async tests."""
-    loop = asyncio.new_event_loop()
-    yield loop
-    loop.close()
-
-
-@pytest.mark.parametrize(
-    "short_id,model_card",
-    TEST_MODELS,
-    ids=[m[0] for m in TEST_MODELS],
-)
-@pytest.mark.asyncio
-async def test_tokenizer_encode_decode(short_id: str, model_card: ModelCard) -> None:
-    """Test that tokenizer can encode and decode text correctly."""
-    model_id = str(model_card.model_id)
-
-    # Download tokenizer files
-    model_path = await download_tokenizer_files(model_id)
-
-    # Verify required files exist
-    has_tokenizer = (
-        (model_path / "tokenizer.json").exists()
-        or (model_path / "tokenizer_config.json").exists()
-        or (model_path / "tiktoken.model").exists()
-        or (model_path / "tokenizer.model").exists()
-    )
-    if not has_tokenizer:
-        pytest.skip(f"Required tokenizer files not found for {model_id}")
-
-    # Load tokenizer
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-
-    # Test basic encoding
-    test_text = "Hello, world!"
-    encoded = tokenizer.encode(test_text)
-    assert isinstance(encoded, list), f"encode() should return a list for {model_id}"
-    assert len(encoded) > 0, f"encode() should return non-empty list for {model_id}"
-    assert all(isinstance(t, int) for t in encoded), (
-        f"All tokens should be integers for {model_id}"
-    )
-
-    # Test decoding
-    decoded = tokenizer.decode(encoded)
-    assert isinstance(decoded, str), f"decode() should return a string for {model_id}"
-    assert test_text in decoded or decoded.strip() == test_text.strip(), (
-        f"decode(encode(x)) should preserve text for {model_id}: got {decoded!r}"
-    )
-
-    # Test with longer text
-    long_text = "The quick brown fox jumps over the lazy dog. " * 10
-    long_encoded = tokenizer.encode(long_text)
-    assert len(long_encoded) > len(encoded), (
-        f"Longer text should produce more tokens for {model_id}"
-    )
-
-    # Test empty string
-    empty_encoded = tokenizer.encode("")
-    assert isinstance(empty_encoded, list), (
-        f"encode('') should return a list for {model_id}"
-    )
-
-    # Test special characters
-    special_text = 'Hello!\n\tWorld? <test> & "quotes"'
-    special_encoded = tokenizer.encode(special_text)
-    assert len(special_encoded) > 0, f"Special chars should encode for {model_id}"
-
-    # Test unicode
-    unicode_text = "Hello 世界 🌍"
-    unicode_encoded = tokenizer.encode(unicode_text)
-    assert len(unicode_encoded) > 0, f"Unicode should encode for {model_id}"
-
-
-@pytest.mark.parametrize(
-    "short_id,model_card",
-    TEST_MODELS,
-    ids=[m[0] for m in TEST_MODELS],
-)
-@pytest.mark.asyncio
-async def test_tokenizer_has_required_attributes(
-    short_id: str, model_card: ModelCard
-) -> None:
-    """Test that tokenizer has required attributes for inference."""
-    model_id = str(model_card.model_id)
-
-    model_path = await download_tokenizer_files(model_id)
-
-    has_tokenizer = (
-        (model_path / "tokenizer.json").exists()
-        or (model_path / "tokenizer_config.json").exists()
-        or (model_path / "tiktoken.model").exists()
-        or (model_path / "tokenizer.model").exists()
-    )
-    if not has_tokenizer:
-        pytest.skip(f"Required tokenizer files not found for {model_id}")
-
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-    eos_token_ids = get_eos_token_ids_for_model(model_id)
-
-    # Check for vocabulary size
-    empty_vocab: dict[str, int] = {}
-    vocab_size: int = getattr(tokenizer, "vocab_size", None) or len(
-        getattr(tokenizer, "get_vocab", lambda: empty_vocab)()
-    )
-    assert vocab_size > 0, f"Tokenizer should have vocab_size > 0 for {model_id}"
-
-    # Check for EOS token (either from tokenizer or explicitly provided)
-    has_eos = (
-        eos_token_ids is not None
-        or getattr(tokenizer, "eos_token_id", None) is not None
-        or getattr(tokenizer, "eos_token", None) is not None
-    )
-    assert has_eos, f"Tokenizer should have EOS token for {model_id}"
-
-
-@pytest.mark.parametrize(
-    "short_id,model_card",
-    TEST_MODELS,
-    ids=[m[0] for m in TEST_MODELS],
-)
-@pytest.mark.asyncio
-async def test_tokenizer_special_tokens(short_id: str, model_card: ModelCard) -> None:
-    """Test that tokenizer can encode text containing special tokens.
-
-    This is critical because the actual inference path uses prompts with
-    special tokens from chat templates. If special tokens aren't handled
-    correctly, encoding will fail.
-    """
-    model_id = str(model_card.model_id)
-
-    model_path = await download_tokenizer_files(model_id)
-
-    has_tokenizer = (
-        (model_path / "tokenizer.json").exists()
-        or (model_path / "tokenizer_config.json").exists()
-        or (model_path / "tiktoken.model").exists()
-        or (model_path / "tokenizer.model").exists()
-    )
-    assert has_tokenizer, f"Required tokenizer files not found for {model_id}"
-
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-
-    # Get special tokens from the tokenizer
-    special_tokens: list[str] = []
-
-    # Try to get special tokens from various sources
-    if hasattr(tokenizer, "all_special_tokens"):
-        special_tokens.extend(tokenizer.all_special_tokens)
-    elif hasattr(tokenizer, "_tokenizer") and hasattr(
-        tokenizer._tokenizer,
-        "all_special_tokens",
-    ):
-        special_tokens.extend(tokenizer._tokenizer.all_special_tokens)
-
-    # Also check for common special token attributes
-    for attr in [
-        "bos_token",
-        "eos_token",
-        "pad_token",
-        "unk_token",
-        "sep_token",
-        "cls_token",
-    ]:
-        token = getattr(tokenizer, attr, None)
-        if token is None and hasattr(tokenizer, "_tokenizer"):
-            token = getattr(tokenizer._tokenizer, attr, None)
-        if token and isinstance(token, str) and token not in special_tokens:
-            special_tokens.append(token)
-
-    # If we found special tokens, test encoding text that contains them
-    if special_tokens:
-        # Create text with special tokens interspersed
-        test_with_special = f"{special_tokens[0]}Hello world"
-        if len(special_tokens) > 1:
-            test_with_special += f"{special_tokens[1]}"
-
-        encoded = tokenizer.encode(test_with_special)
-        assert isinstance(encoded, list), (
-            f"encode() with special tokens should return list for {model_id}"
-        )
-        assert len(encoded) > 0, (
-            f"encode() with special tokens should return non-empty list for {model_id}"
-        )
-        assert all(isinstance(t, int) for t in encoded), (
-            f"All tokens should be integers for {model_id}"
-        )
-
-        # Verify we can decode
-        decoded = tokenizer.decode(encoded)
-        assert isinstance(decoded, str), f"decode() should return string for {model_id}"
-
-    # Test with angle-bracket tokens (common format for special tokens)
-    # These should not raise errors even if they're not actual special tokens
-    angle_bracket_text = "<|test|>Hello<|end|>"
-    encoded = tokenizer.encode(angle_bracket_text)
-    assert isinstance(encoded, list), (
-        f"encode() with angle brackets should return list for {model_id}"
-    )
-    assert len(encoded) > 0, (
-        f"encode() with angle brackets should be non-empty for {model_id}"
-    )
-
-
-# Specifically test Kimi tokenizer since it has special handling
-@pytest.mark.asyncio
-async def test_kimi_tokenizer_specifically():
-    """Test Kimi tokenizer with its specific patches and quirks."""
-    kimi_models = [
-        (short_id, card)
-        for short_id, card in MODEL_CARDS.items()
-        if "kimi" in short_id.lower()
-    ]
-
-    if not kimi_models:
-        pytest.skip("No Kimi models found in MODEL_CARDS")
-
-    _, model_card = kimi_models[0]
-    model_id = str(model_card.model_id)
-
-    model_path = await download_tokenizer_files(model_id)
-
-    # Ensure the custom tokenizer file exists
-    if not (model_path / "tokenization_kimi.py").exists():
-        pytest.skip("tokenization_kimi.py not found")
-
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-    eos_token_ids = get_eos_token_ids_for_model(model_id)
-
-    # Test encode/decode cycle
-    test_text = "Hello, world!"
-    encoded = tokenizer.encode(test_text)
-    decoded = tokenizer.decode(encoded)
-
-    assert len(encoded) > 0, "Kimi tokenizer should encode text"
-    assert isinstance(decoded, str), "Kimi tokenizer should decode to string"
-
-    # Test that the patched encode works (returns list of ints)
-    assert all(isinstance(t, int) for t in encoded), "Tokens should be integers"
-
-    # Test encoding text with special tokens (like from chat templates)
-    # This is critical - the warmup inference uses prompts with special tokens
-    special_token_text = "<|im_user|>user<|im_middle|>Hello<|im_end|><|im_assistant|>"
-    special_encoded = tokenizer.encode(special_token_text)
-    assert len(special_encoded) > 0, "Kimi tokenizer should handle special tokens"
-    assert all(isinstance(t, int) for t in special_encoded), (
-        "Special token encoding should return integers"
-    )
-
-    # Verify EOS token is set
-    assert eos_token_ids == [163586], "Kimi EOS token should be [163586]"
-
-
-# Test GLM tokenizer since it also has special handling
-@pytest.mark.asyncio
-async def test_glm_tokenizer_specifically():
-    """Test GLM tokenizer with its specific EOS tokens."""
-    glm_models = [
-        (short_id, card)
-        for short_id, card in MODEL_CARDS.items()
-        if "glm" in short_id.lower()
-    ]
-
-    if not glm_models:
-        pytest.skip("No GLM models found in MODEL_CARDS")
-
-    _, model_card = glm_models[0]
-    model_id = str(model_card.model_id)
-
-    model_path = await download_tokenizer_files(model_id)
-
-    has_tokenizer = (model_path / "tokenizer.json").exists() or (
-        model_path / "tokenizer_config.json"
-    ).exists()
-    if not has_tokenizer:
-        pytest.skip("GLM tokenizer files not found")
-
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-    eos_token_ids = get_eos_token_ids_for_model(model_id)
-
-    # Test encode/decode
-    test_text = "Hello, world!"
-    encoded = tokenizer.encode(test_text)
-    decoded = tokenizer.decode(encoded)
-
-    assert len(encoded) > 0, "GLM tokenizer should encode text"
-    assert isinstance(decoded, str), "GLM tokenizer should decode to string"
-
-    # Verify EOS tokens
-    assert eos_token_ids == [
-        151336,
-        151329,
-        151338,
-    ], "GLM EOS tokens should be correct"

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

@@ -1,6 +1,5 @@
 import exo.worker.plan as plan_mod
 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
@@ -95,23 +94,13 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
 
     # Local node has already marked its shard as downloaded (not actually used by _load_model)
     local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-        )
+        MODEL_A_ID: DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)
     }
 
     # Global view has completed downloads for both nodes
     global_download_status = {
-        NODE_A: [
-            DownloadCompleted(
-                shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-            )
-        ],
-        NODE_B: [
-            DownloadCompleted(
-                shard_metadata=shard2, node_id=NODE_B, total_bytes=Memory()
-            )
-        ],
+        NODE_A: [DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)],
+        NODE_B: [DownloadCompleted(shard_metadata=shard2, node_id=NODE_B)],
     }
 
     result = plan_mod.plan(
@@ -151,9 +140,7 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
 
     # Local status claims the shard is downloaded already
     local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
-        )
+        MODEL_A_ID: DownloadCompleted(shard_metadata=shard, node_id=NODE_A)
     }
 
     # Global view hasn't caught up yet (no completed shards recorded for NODE_A)
@@ -205,16 +192,10 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
 
     # Only NODE_A's shard is recorded as downloaded globally
     local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-        )
+        MODEL_A_ID: DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)
     }
     global_download_status = {
-        NODE_A: [
-            DownloadCompleted(
-                shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-            )
-        ],
+        NODE_A: [DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)],
         NODE_B: [],  # NODE_B has no downloads completed yet
     }
 
@@ -231,15 +212,9 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     assert result is None
 
     global_download_status = {
-        NODE_A: [
-            DownloadCompleted(
-                shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-            )
-        ],
+        NODE_A: [DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)],
         NODE_B: [
-            DownloadCompleted(
-                shard_metadata=shard2, node_id=NODE_B, total_bytes=Memory()
-            )
+            DownloadCompleted(shard_metadata=shard2, node_id=NODE_B)
         ],  # NODE_B has no downloads completed yet
     }
 

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/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)

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,
@@ -49,12 +50,14 @@ class Tests(BaseModel):
     kind: typing.Literal["init", "warmup", "inference"]
 
 
+hn = socket.gethostname()
 mp.set_start_method("spawn", force=True)
 logger_setup(None)
 
 
 async def main():
     logger.info("starting cool server majig")
+    logger.info(hn)
     await assert_downloads()
     cfg = Config()
     cfg.bind = "0.0.0.0:52415"
@@ -65,6 +68,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,44 +80,32 @@ 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()  # 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))
-    await sd.ensure_shard(
-        await build_full_shard(MODEL_CARDS["llama-3.1-8b-bf16"].model_id)
-    )
-    await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-30b"].model_id))
-    await sd.ensure_shard(
-        await build_full_shard(MODEL_CARDS["gpt-oss-120b-MXFP4-Q8"].model_id)
-    )
-    await sd.ensure_shard(
-        await build_full_shard(MODEL_CARDS["gpt-oss-20b-4bit"].model_id)
-    )
-    await sd.ensure_shard(
-        await build_full_shard(MODEL_CARDS["glm-4.7-8bit-gs32"].model_id)
-    )
-    await sd.ensure_shard(
-        await build_full_shard(MODEL_CARDS["minimax-m2.1-8bit"].model_id)
-    )
+    await sd.ensure_shard(await build_full_shard(MODEL_CARDS["llama-3.2-1b"].model_id))
 
 
 async def ring_backend(test: Tests):
     iid = InstanceId(str(hash(str(test.devs))))
-    weird_hn = socket.gethostname()
-    for dev in test.devs:
-        if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn):
-            hn = dev[0]
-            break
-    else:
-        raise ValueError(f"{weird_hn} not in {test.devs}")
-    return await execute_test(test, ring_instance(test, iid, hn), hn)
+    return await execute_test(test, ring_instance(test, iid))
 
 
-def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
+def ring_instance(test: Tests, iid: InstanceId) -> Instance:
+    global hn
     hbn = [Host(ip="i dont care", port=52416) for _ in test.devs]
     world_size = len(test.devs)
     for i in range(world_size):
-        if test.devs[i][0] == hn:
+        if hn.startswith(test.devs[i][0]):
             hn = test.devs[i][0]
             if i - 1 >= 0:
                 hbn[i - 1] = Host(ip=test.devs[i - 1][1], port=52416)
@@ -121,8 +113,6 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
                 hbn[i + 1] = Host(ip=test.devs[i + 1][1], port=52416)
             hbn[i] = Host(ip="0.0.0.0", port=52416)
             break
-    else:
-        raise ValueError(f"{hn} not in {test.devs}")
 
     meta = MODEL_CARDS[test.model_id].metadata
     instance = MlxRingInstance(
@@ -152,10 +142,10 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
     return instance
 
 
-async def execute_test(test: Tests, instance: Instance, hn: str):
+async def execute_test(test: Tests, instance: Instance):
     world_size = len(test.devs)
     iid = InstanceId(str(hash(str(test.devs))))
-    _handle, recv, send = new_runner(instance, hn)
+    _handle, recv, send = new_runner(instance)
     if world_size > 1:
         send.send(ConnectToGroup(instance_id=iid))
     send.send(LoadModel(instance_id=iid))
@@ -202,19 +192,17 @@ async def execute_test(test: Tests, instance: Instance, hn: str):
 
 async def jaccl_backend(test: Tests):
     iid = InstanceId(str(hash(str(test.devs))))
-    weird_hn = socket.gethostname()
-    for dev in test.devs:
-        if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn):
-            hn = dev[0]
-            break
-    else:
-        raise ValueError(f"{weird_hn} not in {test.devs}")
-    return await execute_test(test, jaccl_instance(test, iid, hn), hn)
+    return await execute_test(test, jaccl_instance(test, iid))
 
 
-def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
+def jaccl_instance(test: Tests, iid: InstanceId):
+    global hn
     meta = MODEL_CARDS[test.model_id].metadata
     world_size = len(test.devs)
+    for name, _ in test.devs:
+        if hn.startswith(name):
+            hn = name
+            break
 
     return MlxJacclInstance(
         instance_id=iid,
@@ -243,7 +231,6 @@ def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
 
 def new_runner(
     instance: Instance,
-    hn: str,
 ) -> tuple[mp.Process, MpReceiver[Event], MpSender[Task]]:
     bound_instance = BoundInstance(
         instance=instance, bound_runner_id=RunnerId(hn), bound_node_id=NodeId(hn)

tests/start_distributed_test.sh

@@ -34,23 +34,19 @@ done
 devs_raw=$(printf "[\"%s\", \"%s\"], " "${weaved[@]}")
 devs="[${devs_raw%, }]"
 
-model_ids=("qwen3-30b" "gpt-oss-120b-MXFP4-Q8" "kimi-k2-thinking")
-
-for model_id in "${model_ids[@]}"; do
-  for i in "${!ips[@]}"; do  
-    { 
-      req="{
-        \"model_id\": \"${model_id}\",
-        \"devs\": ${devs},
-        \"kind\": \"inference\"
-       }"
-      echo "req $req"
-      curl -sN \
-        -X POST "http://${ips[$i]}:52415/${kind}" \
-        -H "Content-Type: application/json" -d "$req" \
-      2>&1 | sed "s/^/\n${hostnames[$i]}@${ips[$i]}: /" || echo "curl to ${hostnames[$i]} failed" && exit 1
-    } &
-  done
-  wait
+for i in "${!ips[@]}"; do  
+  { 
+    req="{
+      \"model_id\": \"llama-3.2-1b\",
+      \"devs\": ${devs},
+      \"kind\": \"inference\"
+     }"
+    echo "req $req"
+    curl -sN \
+      -X POST "http://${ips[$i]}:52415/${kind}" \
+      -H "Content-Type: application/json" -d "$req" \
+    2>&1 | sed "s/^/\n${hostnames[$i]}@${ips[$i]}: /" || echo "curl to ${hostnames[$i]} failed"
+  } &
 done
 
+wait