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/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,96 +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)
-uv run basedpyright
-
-# Linting
-uv run ruff check
-
-# Format code (using nix)
-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
-
-## 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

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

@@ -21,23 +21,21 @@
         "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"
       }
     },
@@ -60,7 +58,7 @@
     "root": {
       "inputs": {
         "fenix": "fenix",
-        "flake-parts": "flake-parts",
+        "flake-utils": "flake-utils",
         "nixpkgs": "nixpkgs",
         "treefmt-nix": "treefmt-nix"
       }
@@ -82,6 +80,21 @@
         "type": "github"
       }
     },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
     "treefmt-nix": {
       "inputs": {
         "nixpkgs": [

flake.nix

@@ -3,131 +3,132 @@
 
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
-
-    flake-parts = {
-      url = "github:hercules-ci/flake-parts";
-      inputs.nixpkgs-lib.follows = "nixpkgs";
-    };
-
+    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";
     };
   };
 
-  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
-      ];
-
-      perSystem =
-        { config, inputs', pkgs, lib, ... }:
-        let
-          fenixToolchain = inputs'.fenix.packages.complete;
-        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 = fenixToolchain.rustfmt;
-              };
-              prettier = {
-                enable = true;
-                includes = [ "*.ts" ];
-              };
-              swift-format.enable = true;
+      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 {
-            packages =
-              [
-                # FORMATTING
-                config.treefmt.build.wrapper
-
-                # PYTHON
-                python313
-                uv
-                ruff
-                basedpyright
-
-                # RUST
-                (fenixToolchain.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"
-              ${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
-            '';
+            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'",

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

@@ -82,7 +82,6 @@ MODEL_CARDS: dict[str, ModelCard] = {
     #         storage_size=Memory.from_kb(754706307),
     #         n_layers=61,
     #         hidden_size=7168,
-    #         supports_tensor=True,
     #     ),
     # ),
     # "deepseek-v3.2-4bit": ModelCard(
@@ -97,7 +96,6 @@ MODEL_CARDS: dict[str, ModelCard] = {
     #         storage_size=Memory.from_kb(754706307 // 2),  # TODO !!!!!
     #         n_layers=61,
     #         hidden_size=7168,
-    #         supports_tensor=True,
     #     ),
     # ),
     # deepseek r1
@@ -556,66 +554,6 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
-    "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-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.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"),

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,10 +1,12 @@
 import http.client
+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
+from exo.shared.types.profiling import NodePerformanceProfile
 
 
 async def check_reachability(
@@ -58,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,23 +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.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):
@@ -96,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:
@@ -103,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] :]
@@ -134,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,
@@ -149,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
@@ -163,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), (
@@ -198,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,
@@ -243,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,
@@ -252,7 +231,7 @@ def tensor_auto_parallel(
             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,
@@ -260,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)}")
 
@@ -314,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
                 )
@@ -360,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)
@@ -408,9 +353,7 @@ 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)
@@ -438,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/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/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