add GitHub releases with Sparkle release notes integration

- Require draft GitHub release with markdown notes before pushing tag
- Upgrade to Sparkle 2.9.0-beta.1 with native markdown support
- Inject release notes directly into appcast.xml
- Publish draft release with DMG attached after build

Closes #1140

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

.github/workflows/build-app.yml

@@ -22,8 +22,6 @@ on:
 jobs:
   build-macos-app:
     runs-on: "macos-26"
-    permissions:
-      contents: write
     env:
       SPARKLE_VERSION: 2.9.0-beta.1
       SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }}
@@ -105,11 +103,10 @@ jobs:
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
-          # Find draft release by name using gh release list (more reliable with default token)
-          echo "Looking for draft release named '$GITHUB_REF_NAME'..."
-          DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "")
+          # Find draft release - match by name (title) since draft releases don't have tag_name set yet
+          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq "[.[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")] | first // empty" 2>/dev/null || echo "")
 
-          if [[ -z "$DRAFT_EXISTS" ]]; then
+          if [[ -z "$RELEASE_JSON" ]]; then
             if [[ "$IS_ALPHA" == "true" ]]; then
               echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)"
               echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
@@ -120,10 +117,6 @@ jobs:
             exit 1
           fi
 
-          # Fetch full release details via API to get body and ID
-          echo "Found draft release, fetching details..."
-          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "")
-
           # Extract release notes
           NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""')
           if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then

AGENTS.md

@@ -40,31 +40,6 @@ uv run ruff check
 nix fmt
 ```
 
-## Pre-Commit Checks (REQUIRED)
-
-**IMPORTANT: Always run these checks before committing code. CI will fail if these don't pass.**
-
-```bash
-# 1. Type checking - MUST pass with 0 errors
-uv run basedpyright
-
-# 2. Linting - MUST pass
-uv run ruff check
-
-# 3. Formatting - MUST be applied
-nix fmt
-
-# 4. Tests - MUST pass
-uv run pytest
-```
-
-Run all checks in sequence:
-```bash
-uv run basedpyright && uv run ruff check && nix fmt && uv run pytest
-```
-
-If `nix fmt` changes any files, stage them before committing. The CI runs `nix flake check` which verifies formatting, linting, and runs Rust tests.
-
 ## Architecture
 
 ### Node Composition

README.md

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

bench/exo_bench.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import argparse
-import contextlib
 import http.client
 import json
 import os
@@ -27,7 +26,7 @@ class ExoHttpError(RuntimeError):
 
 
 class ExoClient:
-    def __init__(self, host: str, port: int, timeout_s: float = 600.0):
+    def __init__(self, host: str, port: int, timeout_s: float = 2400.0):
         self.host = host
         self.port = port
         self.timeout_s = timeout_s
@@ -105,46 +104,22 @@ def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
 
-def runner_failed(runner: dict[str, Any]) -> bool:
-    return "RunnerFailed" in runner
-
-
-def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
-    if "RunnerFailed" in runner:
-        return runner["RunnerFailed"].get("errorMessage")
-    return None
-
-
 def wait_for_instance_ready(
     client: ExoClient, instance_id: str, timeout: float = 24000.0
 ) -> None:
     start_time = time.time()
-    instance_existed = False
     while time.time() - start_time < timeout:
         state = client.request_json("GET", "/state")
         instances = state.get("instances", {})
 
         if instance_id not in instances:
-            if instance_existed:
-                # Instance was deleted after being created - likely due to runner failure
-                raise RuntimeError(
-                    f"Instance {instance_id} was deleted (runner may have failed)"
-                )
             time.sleep(0.1)
             continue
 
-        instance_existed = True
         instance = instances[instance_id]
         runner_ids = runner_ids_from_instance(instance)
         runners = state.get("runners", {})
 
-        # Check for failed runners first
-        for rid in runner_ids:
-            runner = runners.get(rid, {})
-            if runner_failed(runner):
-                error_msg = get_runner_failed_message(runner) or "Unknown error"
-                raise RuntimeError(f"Runner {rid} failed: {error_msg}")
-
         if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
             return
 
@@ -324,12 +299,6 @@ def main() -> int:
         default=4,
         help="Only consider placements using <= this many nodes.",
     )
-    ap.add_argument(
-        "--min-nodes",
-        type=int,
-        default=1,
-        help="Only consider placements using >= this many nodes.",
-    )
     ap.add_argument(
         "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
     )
@@ -351,7 +320,7 @@ def main() -> int:
         help="Warmup runs per placement (uses first pp/tg).",
     )
     ap.add_argument(
-        "--timeout", type=float, default=600.0, help="HTTP timeout (seconds)."
+        "--timeout", type=float, default=2400.0, help="HTTP timeout (seconds)."
     )
     ap.add_argument(
         "--json-out",
@@ -430,7 +399,7 @@ def main() -> int:
         ):
             continue
 
-        if args.min_nodes <= n <= args.max_nodes:
+        if 0 < n <= args.max_nodes:
             selected.append(p)
 
     if not selected:
@@ -472,13 +441,7 @@ def main() -> int:
         )
 
         client.request_json("POST", "/instance", body={"instance": instance})
-        try:
-            wait_for_instance_ready(client, instance_id)
-        except (RuntimeError, TimeoutError) as e:
-            logger.error(f"Failed to initialize placement: {e}")
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{instance_id}")
-            continue
+        wait_for_instance_ready(client, instance_id)
 
         time.sleep(1)
 

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

@@ -71,36 +71,35 @@ export interface Instance {
 	};
 }
 
-// Split state interfaces
-interface RawNodeIdentity {
-	modelId: string;
-	chipId: string;
-	friendlyName: string;
-}
-
-interface RawNodeMemory {
-	ramTotal: { inBytes: number };
-	ramAvailable: { inBytes: number };
-	swapTotal: { inBytes: number };
-	swapAvailable: { inBytes: number };
-}
-
-interface RawNodeSystem {
-	gpuUsage?: number;
-	temp?: number;
-	sysPower?: number;
-	pcpuUsage?: number;
-	ecpuUsage?: number;
-	anePower?: number;
-}
-
-interface RawNetworkInterface {
-	name: string;
-	ipAddress: string;
+interface RawNodeProfile {
+	modelId?: string;
+	chipId?: string;
+	friendlyName?: string;
+	networkInterfaces?: Array<{
+		name?: string;
+		ipAddress?: string;
+		addresses?: Array<{ address?: string } | string>;
+		ipv4?: string;
+		ipv6?: string;
+		ipAddresses?: string[];
+		ips?: string[];
+	}>;
+	memory?: {
+		ramTotal?: { inBytes: number };
+		ramAvailable?: { inBytes: number };
+		swapTotal?: { inBytes: number };
+		swapAvailable?: { inBytes: number };
+	};
+	system?: {
+		gpuUsage?: number;
+		temp?: number;
+		sysPower?: number;
+	};
 }
 
 interface RawTopologyNode {
 	nodeId: string;
+	nodeProfile: RawNodeProfile;
 }
 
 interface RawTopologyConnection {
@@ -116,6 +115,8 @@ interface RawTopology {
 	connections?: RawTopologyConnection[];
 }
 
+type RawNodeProfiles = Record<string, RawNodeProfile>;
+
 export interface DownloadProgress {
 	totalBytes: number;
 	downloadedBytes: number;
@@ -170,11 +171,7 @@ interface RawStateResponse {
 	>;
 	runners?: Record<string, unknown>;
 	downloads?: Record<string, unknown[]>;
-	// Split state fields
-	nodeIdentities?: Record<string, RawNodeIdentity>;
-	nodeMemories?: Record<string, RawNodeMemory>;
-	nodeSystems?: Record<string, RawNodeSystem>;
-	nodeNetworks?: Record<string, RawNetworkInterface[]>;
+	nodeProfiles?: RawNodeProfiles;
 }
 
 export interface MessageAttachment {
@@ -211,41 +208,66 @@ const STORAGE_KEY = "exo-conversations";
 
 function transformTopology(
 	raw: RawTopology,
-	identities?: Record<string, RawNodeIdentity>,
-	memories?: Record<string, RawNodeMemory>,
-	systems?: Record<string, RawNodeSystem>,
-	networks?: Record<string, RawNetworkInterface[]>,
+	profiles?: RawNodeProfiles,
 ): TopologyData {
 	const nodes: Record<string, NodeInfo> = {};
 	const edges: TopologyEdge[] = [];
 
 	for (const node of raw.nodes || []) {
-		// Get split state fields (may be undefined if events haven't arrived yet)
-		const identity = identities?.[node.nodeId];
-		const memory = memories?.[node.nodeId];
-		const system = systems?.[node.nodeId];
-		const network = networks?.[node.nodeId];
-
-		const ramTotal = memory?.ramTotal?.inBytes ?? 0;
-		const ramAvailable = memory?.ramAvailable?.inBytes ?? 0;
+		const mergedProfile = profiles?.[node.nodeId];
+		const profile = { ...(node.nodeProfile ?? {}), ...(mergedProfile ?? {}) };
+		const ramTotal = profile?.memory?.ramTotal?.inBytes ?? 0;
+		const ramAvailable = profile?.memory?.ramAvailable?.inBytes ?? 0;
 		const ramUsage = Math.max(ramTotal - ramAvailable, 0);
 
-		const networkInterfaces = (network ?? []).map((iface) => ({
-			name: iface.name,
-			addresses: [iface.ipAddress],
-		}));
+		const networkInterfaces = (profile?.networkInterfaces || []).map(
+			(iface) => {
+				const addresses: string[] = [];
+				if (iface.ipAddress && typeof iface.ipAddress === "string") {
+					addresses.push(iface.ipAddress);
+				}
+				if (Array.isArray(iface.addresses)) {
+					for (const addr of iface.addresses) {
+						if (typeof addr === "string") addresses.push(addr);
+						else if (addr && typeof addr === "object" && addr.address)
+							addresses.push(addr.address);
+					}
+				}
+				if (Array.isArray(iface.ipAddresses)) {
+					addresses.push(
+						...iface.ipAddresses.filter(
+							(a): a is string => typeof a === "string",
+						),
+					);
+				}
+				if (Array.isArray(iface.ips)) {
+					addresses.push(
+						...iface.ips.filter((a): a is string => typeof a === "string"),
+					);
+				}
+				if (iface.ipv4 && typeof iface.ipv4 === "string")
+					addresses.push(iface.ipv4);
+				if (iface.ipv6 && typeof iface.ipv6 === "string")
+					addresses.push(iface.ipv6);
+
+				return {
+					name: iface.name,
+					addresses: Array.from(new Set(addresses)),
+				};
+			},
+		);
 
 		const ipToInterface: Record<string, string> = {};
 		for (const iface of networkInterfaces) {
-			for (const addr of iface.addresses) {
-				ipToInterface[addr] = iface.name;
+			for (const addr of iface.addresses || []) {
+				ipToInterface[addr] = iface.name ?? "";
 			}
 		}
 
 		nodes[node.nodeId] = {
 			system_info: {
-				model_id: identity?.modelId ?? "Unknown",
-				chip: identity?.chipId,
+				model_id: profile?.modelId ?? "Unknown",
+				chip: profile?.chipId,
 				memory: ramTotal,
 			},
 			network_interfaces: networkInterfaces,
@@ -256,15 +278,17 @@ function transformTopology(
 					ram_total: ramTotal,
 				},
 				temp:
-					system?.temp !== undefined
-						? { gpu_temp_avg: system.temp }
+					profile?.system?.temp !== undefined
+						? { gpu_temp_avg: profile.system.temp }
 						: undefined,
 				gpu_usage:
-					system?.gpuUsage !== undefined ? [0, system.gpuUsage] : undefined,
-				sys_power: system?.sysPower,
+					profile?.system?.gpuUsage !== undefined
+						? [0, profile.system.gpuUsage]
+						: undefined,
+				sys_power: profile?.system?.sysPower,
 			},
 			last_macmon_update: Date.now() / 1000,
-			friendly_name: identity?.friendlyName,
+			friendly_name: profile?.friendlyName,
 		};
 	}
 
@@ -844,13 +868,7 @@ class AppStore {
 			const data: RawStateResponse = await response.json();
 
 			if (data.topology) {
-				this.topologyData = transformTopology(
-					data.topology,
-					data.nodeIdentities,
-					data.nodeMemories,
-					data.nodeSystems,
-					data.nodeNetworks,
-				);
+				this.topologyData = transformTopology(data.topology, data.nodeProfiles);
 			}
 			if (data.instances) {
 				this.instances = data.instances;

pyproject.toml

@@ -126,6 +126,3 @@ env = [
   "EXO_TESTS=1"
 ]
 addopts = "-m 'not slow'"
-filterwarnings = [
-    "ignore:builtin type Swig:DeprecationWarning",
-]

src/exo/main.py

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

src/exo/master/api.py

@@ -1,14 +1,13 @@
 import time
 from collections.abc import AsyncGenerator
-from http import HTTPStatus
 from typing import cast
 
 import anyio
-from anyio import BrokenResourceError, create_task_group
+from anyio import create_task_group
 from anyio.abc import TaskGroup
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType]
 from hypercorn.config import Config
@@ -30,8 +29,6 @@ from exo.shared.types.api import (
     CreateInstanceParams,
     CreateInstanceResponse,
     DeleteInstanceResponse,
-    ErrorInfo,
-    ErrorResponse,
     FinishReason,
     GenerationStats,
     ModelList,
@@ -52,12 +49,7 @@ from exo.shared.types.commands import (
     TaskFinished,
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId
-from exo.shared.types.events import (
-    ChunkGenerated,
-    Event,
-    ForwarderEvent,
-    IndexedEvent,
-)
+from exo.shared.types.events import ChunkGenerated, Event, ForwarderEvent, IndexedEvent
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.state import State
@@ -123,7 +115,6 @@ class API:
         self.paused_ev: anyio.Event = anyio.Event()
 
         self.app = FastAPI()
-        self._setup_exception_handlers()
         self._setup_cors()
         self._setup_routes()
 
@@ -154,21 +145,6 @@ class API:
         self.paused_ev.set()
         self.paused_ev = anyio.Event()
 
-    def _setup_exception_handlers(self) -> None:
-        self.app.exception_handler(HTTPException)(self.http_exception_handler)
-
-    async def http_exception_handler(
-        self, _: Request, exc: HTTPException
-    ) -> JSONResponse:
-        err = ErrorResponse(
-            error=ErrorInfo(
-                message=exc.detail,
-                type=HTTPStatus(exc.status_code).phrase,
-                code=exc.status_code,
-            )
-        )
-        return JSONResponse(err.model_dump(), status_code=exc.status_code)
-
     def _setup_cors(self) -> None:
         self.app.add_middleware(
             CORSMiddleware,
@@ -430,18 +406,6 @@ class API:
         """Generate chat completion stream as JSON strings."""
 
         async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                error_response = ErrorResponse(
-                    error=ErrorInfo(
-                        message=chunk.error_message or "Internal server error",
-                        type="InternalServerError",
-                        code=500,
-                    )
-                )
-                yield f"data: {error_response.model_dump_json()}\n\n"
-                yield "data: [DONE]\n\n"
-                return
-
             chunk_response: ChatCompletionResponse = chunk_to_response(
                 chunk, command_id
             )
@@ -462,12 +426,6 @@ class API:
         finish_reason: FinishReason | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                raise HTTPException(
-                    status_code=500,
-                    detail=chunk.error_message or "Internal server error",
-                )
-
             if model is None:
                 model = chunk.model
 
@@ -505,12 +463,6 @@ class API:
         stats: GenerationStats | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                raise HTTPException(
-                    status_code=500,
-                    detail=chunk.error_message or "Internal server error",
-                )
-
             if model is None:
                 model = chunk.model
 
@@ -600,8 +552,9 @@ class API:
         """Calculate total available memory across all nodes in bytes."""
         total_available = Memory()
 
-        for memory in self.state.node_memories.values():
-            total_available += memory.ram_available
+        for node in self.state.topology.list_nodes():
+            if node.node_profile is not None:
+                total_available += node.node_profile.memory.ram_available
 
         return total_available
 
@@ -654,14 +607,14 @@ class API:
                 for idx, event in self.event_buffer.drain_indexed():
                     self._event_log.append(event)
                     self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
-                    if isinstance(event, ChunkGenerated):
+                    if (
+                        isinstance(event, ChunkGenerated)
+                        and event.command_id in self._chat_completion_queues
+                    ):
                         assert isinstance(event.chunk, TokenChunk)
-                        queue = self._chat_completion_queues.get(event.command_id)
-                        if queue is not None:
-                            try:
-                                await queue.send(event.chunk)
-                            except BrokenResourceError:
-                                self._chat_completion_queues.pop(event.command_id, None)
+                        await self._chat_completion_queues[event.command_id].send(
+                            event.chunk
+                        )
 
     async def _pause_on_new_election(self):
         with self.election_receiver as ems:

src/exo/master/placement.py

@@ -113,7 +113,6 @@ def place_instance(
                 node.node_profile.memory.ram_available
                 for node in cycle
                 if node.node_profile is not None
-                and node.node_profile.memory is not None
             ),
             start=Memory(),
         ),

src/exo/master/placement_utils.py

@@ -25,10 +25,7 @@ class NodeWithProfile(BaseModel):
 
 
 def narrow_all_nodes(nodes: list[NodeInfo]) -> TypeGuard[list[NodeWithProfile]]:
-    return all(
-        node.node_profile is not None and node.node_profile.memory is not None
-        for node in nodes
-    )
+    return all(node.node_profile is not None for node in nodes)
 
 
 def filter_cycles_by_memory(
@@ -39,14 +36,8 @@ def filter_cycles_by_memory(
         if not narrow_all_nodes(cycle):
             continue
 
-        # narrow_all_nodes guarantees memory is not None
         total_mem = sum(
-            (
-                node.node_profile.memory.ram_available
-                for node in cycle
-                if node.node_profile.memory is not None
-            ),
-            start=Memory(),
+            (node.node_profile.memory.ram_available for node in cycle), start=Memory()
         )
         if total_mem >= required_memory:
             filtered_cycles.append(cast(list[NodeInfo], cycle))
@@ -58,89 +49,33 @@ def get_smallest_cycles(cycles: list[list[NodeInfo]]) -> list[list[NodeInfo]]:
     return [cycle for cycle in cycles if len(cycle) == min_nodes]
 
 
-def allocate_layers_proportionally(
-    total_layers: int,
-    memory_fractions: list[float],
-) -> list[int]:
-    n = len(memory_fractions)
-    if n == 0:
-        raise ValueError("Cannot allocate layers to an empty node list")
-    if total_layers < n:
-        raise ValueError(
-            f"Cannot distribute {total_layers} layers across {n} nodes "
-            "(need at least 1 layer per node)"
-        )
-
-    # Largest remainder: floor each, then distribute remainder by fractional part
-    raw = [f * total_layers for f in memory_fractions]
-    result = [int(r) for r in raw]
-    by_remainder = sorted(range(n), key=lambda i: raw[i] - result[i], reverse=True)
-    for i in range(total_layers - sum(result)):
-        result[by_remainder[i]] += 1
-
-    # Ensure minimum 1 per node by taking from the largest
-    for i in range(n):
-        if result[i] == 0:
-            max_idx = max(range(n), key=lambda j: result[j])
-            assert result[max_idx] > 1
-            result[max_idx] -= 1
-            result[i] = 1
-
-    return result
-
-
 def get_shard_assignments_for_pipeline_parallel(
     model_meta: ModelMetadata,
     selected_cycle: list[NodeWithProfile],
 ):
-    if not selected_cycle:
-        raise ValueError("Cannot create shard assignments for empty node cycle")
-
     cycle_memory = sum(
-        (
-            node.node_profile.memory.ram_available
-            for node in selected_cycle
-            if node.node_profile.memory is not None
-        ),
+        (node.node_profile.memory.ram_available for node in selected_cycle),
         start=Memory(),
     )
-
-    if cycle_memory.in_bytes == 0:
-        raise ValueError("Cannot create shard assignments: total available memory is 0")
-
     total_layers = model_meta.n_layers
     world_size = len(selected_cycle)
     runner_to_shard: dict[RunnerId, ShardMetadata] = {}
     node_to_runner: dict[NodeId, RunnerId] = {}
 
-    layer_allocations = allocate_layers_proportionally(
-        total_layers=total_layers,
-        memory_fractions=[
-            node.node_profile.memory.ram_available.in_bytes / cycle_memory.in_bytes
-            for node in selected_cycle
-            if node.node_profile.memory is not None
-        ],
-    )
-
-    # Validate each node has sufficient memory for its assigned layers
-    memory_per_layer = model_meta.storage_size.in_bytes / total_layers
-    for i, (node, node_layers) in enumerate(
-        zip(selected_cycle, layer_allocations, strict=True)
-    ):
-        assert node.node_profile.memory is not None
-        required_memory = node_layers * memory_per_layer
-        available_memory = node.node_profile.memory.ram_available.in_bytes
-        if required_memory > available_memory:
-            raise ValueError(
-                f"Node {i} ({node.node_id}) has insufficient memory: "
-                f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
-                f"but only has {available_memory / (1024**3):.2f} GB available"
-            )
-
     layers_assigned = 0
-    for i, (node, node_layers) in enumerate(
-        zip(selected_cycle, layer_allocations, strict=True)
-    ):
+    for i, node in enumerate(selected_cycle):
+        if i == len(selected_cycle) - 1:
+            node_layers = total_layers - layers_assigned
+        else:
+            node_layers = round(
+                total_layers
+                * (
+                    node.node_profile.memory.ram_available.in_bytes
+                    / cycle_memory.in_bytes
+                )
+            )
+            node_layers = max(1, node_layers)
+
         runner_id = RunnerId()
 
         shard = PipelineShardMetadata(

src/exo/master/tests/test_api_error_handling.py

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

src/exo/master/tests/test_master.py

@@ -19,13 +19,16 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InstanceCreated,
-    NodeIdentityMeasured,
-    NodeMemoryMeasured,
+    NodePerformanceMeasured,
     TaskCreated,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.profiling import MemoryPerformanceProfile
+from exo.shared.types.profiling import (
+    MemoryPerformanceProfile,
+    NodePerformanceProfile,
+    SystemPerformanceProfile,
+)
 from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
 from exo.shared.types.tasks import TaskStatus
 from exo.shared.types.worker.instances import (
@@ -72,39 +75,29 @@ async def test_master():
         tg.start_soon(master.run)
 
         sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
-        # inject NodeIdentityMeasured and NodeMemoryMeasured events
-        logger.info("inject NodeIdentityMeasured event")
+        # inject a NodePerformanceProfile event
+        logger.info("inject a NodePerformanceProfile event")
         await local_event_sender.send(
             ForwarderEvent(
                 origin_idx=0,
                 origin=sender_node_id,
                 session=session_id,
                 event=(
-                    NodeIdentityMeasured(
+                    NodePerformanceMeasured(
                         when=str(datetime.now(tz=timezone.utc)),
                         node_id=node_id,
-                        model_id="maccy",
-                        chip_id="arm",
-                        friendly_name="test",
-                    )
-                ),
-            )
-        )
-        logger.info("inject NodeMemoryMeasured event")
-        await local_event_sender.send(
-            ForwarderEvent(
-                origin_idx=1,
-                origin=sender_node_id,
-                session=session_id,
-                event=(
-                    NodeMemoryMeasured(
-                        when=str(datetime.now(tz=timezone.utc)),
-                        node_id=node_id,
-                        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),
+                        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(),
                         ),
                     )
                 ),
@@ -115,7 +108,7 @@ async def test_master():
         logger.info("wait for initial topology event")
         while len(list(master.state.topology.list_nodes())) == 0:
             await anyio.sleep(0.001)
-        while len(master.state.node_identities) == 0:
+        while len(master.state.node_profiles) == 0:
             await anyio.sleep(0.001)
 
         logger.info("inject a CreateInstance Command")
@@ -162,19 +155,17 @@ async def test_master():
                 ),
             )
         )
-        while len(_get_events()) < 4:
+        while len(_get_events()) < 3:
             await anyio.sleep(0.01)
 
         events = _get_events()
-        assert len(events) == 4
+        assert len(events) == 3
         assert events[0].idx == 0
         assert events[1].idx == 1
         assert events[2].idx == 2
-        assert events[3].idx == 3
-        assert isinstance(events[0].event, NodeIdentityMeasured)
-        assert isinstance(events[1].event, NodeMemoryMeasured)
-        assert isinstance(events[2].event, InstanceCreated)
-        created_instance = events[2].event.instance
+        assert isinstance(events[0].event, NodePerformanceMeasured)
+        assert isinstance(events[1].event, InstanceCreated)
+        created_instance = events[1].event.instance
         assert isinstance(created_instance, MlxRingInstance)
         runner_id = list(created_instance.shard_assignments.runner_to_shard.keys())[0]
         # Validate the shard assignments
@@ -206,10 +197,10 @@ async def test_master():
         assert len(created_instance.hosts_by_node[node_id]) == 1
         assert created_instance.hosts_by_node[node_id][0].ip == "0.0.0.0"
         assert created_instance.ephemeral_port > 0
-        assert isinstance(events[3].event, TaskCreated)
-        assert events[3].event.task.task_status == TaskStatus.Pending
-        assert isinstance(events[3].event.task, ChatCompletionTask)
-        assert events[3].event.task.task_params == ChatCompletionTaskParams(
+        assert isinstance(events[2].event, TaskCreated)
+        assert events[2].event.task.task_status == TaskStatus.Pending
+        assert isinstance(events[2].event.task, ChatCompletionTask)
+        assert events[2].event.task.task_params == ChatCompletionTaskParams(
             model="llama-3.2-1b",
             messages=[
                 ChatCompletionMessage(role="user", content="Hello, how are you?")

src/exo/master/tests/test_placement.py

@@ -70,7 +70,7 @@ def place_instance_command(model_meta: ModelMetadata) -> PlaceInstance:
     [
         ((500, 500, 1000), 12, (3, 3, 6)),
         ((500, 500, 500), 12, (4, 4, 4)),
-        ((312, 468, 1092), 12, (2, 3, 7)),
+        ((312, 518, 1024), 12, (2, 3, 7)),
     ],
 )
 def test_get_instance_placements_create_instance(

src/exo/master/tests/test_placement_utils.py

@@ -3,7 +3,6 @@ from typing import Callable
 import pytest
 
 from exo.master.placement_utils import (
-    allocate_layers_proportionally,
     filter_cycles_by_memory,
     get_hosts_from_subgraph,
     get_mlx_jaccl_coordinators,
@@ -166,9 +165,6 @@ def test_get_smallest_cycles(
         ((500, 500, 1000), 12, (3, 3, 6)),
         ((500, 500, 500), 12, (4, 4, 4)),
         ((312, 518, 1024), 12, (2, 3, 7)),
-        # Edge case: one node has ~90% of memory - should not over-allocate.
-        # Each node must have enough memory for at least 1 layer (50 KB = 1000/20).
-        ((900, 50, 50), 20, (18, 1, 1)),
     ],
 )
 def test_get_shard_assignments(
@@ -401,96 +397,3 @@ def test_get_mlx_jaccl_coordinators(
     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"
-
-
-class TestAllocateLayersProportionally:
-    def test_empty_node_list_raises(self):
-        with pytest.raises(ValueError, match="empty node list"):
-            allocate_layers_proportionally(total_layers=10, memory_fractions=[])
-
-    def test_zero_layers_raises(self):
-        with pytest.raises(ValueError, match="need at least 1 layer per node"):
-            allocate_layers_proportionally(total_layers=0, memory_fractions=[0.5, 0.5])
-
-    def test_negative_layers_raises(self):
-        with pytest.raises(ValueError, match="need at least 1 layer per node"):
-            allocate_layers_proportionally(total_layers=-1, memory_fractions=[0.5, 0.5])
-
-    def test_fewer_layers_than_nodes_raises(self):
-        with pytest.raises(ValueError, match="need at least 1 layer per node"):
-            allocate_layers_proportionally(
-                total_layers=2, memory_fractions=[0.33, 0.33, 0.34]
-            )
-
-    def test_equal_distribution(self):
-        result = allocate_layers_proportionally(
-            total_layers=12, memory_fractions=[0.25, 0.25, 0.25, 0.25]
-        )
-        assert result == [3, 3, 3, 3]
-        assert sum(result) == 12
-
-    def test_proportional_distribution(self):
-        result = allocate_layers_proportionally(
-            total_layers=12, memory_fractions=[0.25, 0.25, 0.50]
-        )
-        assert result == [3, 3, 6]
-        assert sum(result) == 12
-
-    def test_extreme_imbalance_ensures_minimum(self):
-        result = allocate_layers_proportionally(
-            total_layers=20, memory_fractions=[0.975, 0.0125, 0.0125]
-        )
-        assert all(layers >= 1 for layers in result)
-        assert sum(result) == 20
-        # Small nodes get minimum 1 layer
-        assert result == [18, 1, 1]
-
-    def test_single_node_gets_all_layers(self):
-        result = allocate_layers_proportionally(total_layers=10, memory_fractions=[1.0])
-        assert result == [10]
-
-    def test_minimum_viable_allocation(self):
-        result = allocate_layers_proportionally(
-            total_layers=3, memory_fractions=[0.33, 0.33, 0.34]
-        )
-        assert result == [1, 1, 1]
-        assert sum(result) == 3
-
-
-def test_get_shard_assignments_insufficient_memory_raises(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
-    """Test that ValueError is raised when a node has insufficient memory for its layers."""
-    node_a_id = NodeId()
-    node_b_id = NodeId()
-    node_c_id = NodeId()
-
-    # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers)
-    node_a = create_node(900 * 1024, node_a_id)
-    node_b = create_node(50 * 1024, node_b_id)
-    node_c = create_node(10 * 1024, node_c_id)  # Insufficient memory
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
-
-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
-
-    model_meta = ModelMetadata(
-        model_id=ModelId("test-model"),
-        pretty_name="Test Model",
-        n_layers=20,
-        storage_size=Memory.from_kb(1000),
-        hidden_size=1000,
-        supports_tensor=True,
-    )
-    cycles = topology.get_cycles()
-    selected_cycle = cycles[0]
-
-    with pytest.raises(ValueError, match="insufficient memory"):
-        get_shard_assignments(model_meta, selected_cycle, Sharding.Pipeline)

src/exo/shared/apply.py

@@ -13,10 +13,8 @@ from exo.shared.types.events import (
     InstanceDeleted,
     NodeCreated,
     NodeDownloadProgress,
-    NodeIdentityMeasured,
     NodeMemoryMeasured,
-    NodeNetworkMeasured,
-    NodeSystemMeasured,
+    NodePerformanceMeasured,
     NodeTimedOut,
     RunnerDeleted,
     RunnerStatusUpdated,
@@ -29,13 +27,7 @@ from exo.shared.types.events import (
     TopologyEdgeCreated,
     TopologyEdgeDeleted,
 )
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NetworkInterfaceInfo,
-    NodeIdentity,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
-)
+from exo.shared.types.profiling import NodePerformanceProfile, SystemPerformanceProfile
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import NodeInfo
@@ -59,12 +51,8 @@ def event_apply(event: Event, state: State) -> State:
             return apply_topology_node_created(event, state)
         case NodeTimedOut():
             return apply_node_timed_out(event, state)
-        case NodeIdentityMeasured():
-            return apply_node_identity_measured(event, state)
-        case NodeSystemMeasured():
-            return apply_node_system_measured(event, state)
-        case NodeNetworkMeasured():
-            return apply_node_network_measured(event, state)
+        case NodePerformanceMeasured():
+            return apply_node_performance_measured(event, state)
         case NodeDownloadProgress():
             return apply_node_download_progress(event, state)
         case NodeMemoryMeasured():
@@ -202,19 +190,8 @@ def apply_runner_deleted(event: RunnerDeleted, state: State) -> State:
 def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     topology = copy.copy(state.topology)
     state.topology.remove_node(event.node_id)
-    node_identities = {
-        key: value
-        for key, value in state.node_identities.items()
-        if key != event.node_id
-    }
-    node_memories = {
-        key: value for key, value in state.node_memories.items() if key != event.node_id
-    }
-    node_systems = {
-        key: value for key, value in state.node_systems.items() if key != event.node_id
-    }
-    node_networks = {
-        key: value for key, value in state.node_networks.items() if key != event.node_id
+    node_profiles = {
+        key: value for key, value in state.node_profiles.items() if key != event.node_id
     }
     last_seen = {
         key: value for key, value in state.last_seen.items() if key != event.node_id
@@ -222,120 +199,32 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     return state.model_copy(
         update={
             "topology": topology,
-            "node_identities": node_identities,
-            "node_memories": node_memories,
-            "node_systems": node_systems,
-            "node_networks": node_networks,
+            "node_profiles": node_profiles,
             "last_seen": last_seen,
         }
     )
 
 
-def _reconstruct_profile(
-    node_id: NodeId,
-    state: State,
-    *,
-    identity: NodeIdentity | None = None,
-    memory: MemoryPerformanceProfile | None = None,
-    system: SystemPerformanceProfile | None = None,
-    network_interfaces: list[NetworkInterfaceInfo] | None = None,
-) -> NodePerformanceProfile:
-    """Reconstruct a NodePerformanceProfile from split state storage.
-
-    Uses provided overrides, falling back to state values.
-    """
-    ident = identity or state.node_identities.get(node_id)
-    mem = memory or state.node_memories.get(node_id)
-    sys = system or state.node_systems.get(node_id)
-    nets = (
-        network_interfaces
-        if network_interfaces is not None
-        else state.node_networks.get(node_id, [])
-    )
-
-    return NodePerformanceProfile(
-        model_id=ident.model_id if ident else None,
-        chip_id=ident.chip_id if ident else None,
-        friendly_name=ident.friendly_name if ident else None,
-        memory=mem,
-        network_interfaces=nets,
-        system=sys,
-    )
-
-
-def apply_node_identity_measured(event: NodeIdentityMeasured, state: State) -> State:
-    topology = copy.copy(state.topology)
-
-    identity = NodeIdentity(
-        model_id=event.model_id,
-        chip_id=event.chip_id,
-        friendly_name=event.friendly_name,
-    )
-    new_identities: Mapping[NodeId, NodeIdentity] = {
-        **state.node_identities,
-        event.node_id: identity,
+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),
     }
-    if not topology.contains_node(event.node_id):
-        topology.add_node(NodeInfo(node_id=event.node_id))
-    reconstructed = _reconstruct_profile(event.node_id, state, identity=identity)
-    topology.update_node_profile(event.node_id, reconstructed)
-    return state.model_copy(
-        update={
-            "node_identities": new_identities,
-            "topology": topology,
-            "last_seen": last_seen,
-        }
-    )
-
-
-def apply_node_system_measured(event: NodeSystemMeasured, state: State) -> State:
+    state = state.model_copy(update={"node_profiles": new_profiles})
     topology = copy.copy(state.topology)
-
-    new_systems: Mapping[NodeId, SystemPerformanceProfile] = {
-        **state.node_systems,
-        event.node_id: event.system,
-    }
-    last_seen: Mapping[NodeId, datetime] = {
-        **state.last_seen,
-        event.node_id: datetime.fromisoformat(event.when),
-    }
+    # TODO: NodeCreated
     if not topology.contains_node(event.node_id):
         topology.add_node(NodeInfo(node_id=event.node_id))
-    reconstructed = _reconstruct_profile(event.node_id, state, system=event.system)
-    topology.update_node_profile(event.node_id, reconstructed)
+    topology.update_node_profile(event.node_id, event.node_profile)
     return state.model_copy(
         update={
-            "node_systems": new_systems,
-            "topology": topology,
-            "last_seen": last_seen,
-        }
-    )
-
-
-def apply_node_network_measured(event: NodeNetworkMeasured, state: State) -> State:
-    topology = copy.copy(state.topology)
-
-    new_networks: Mapping[NodeId, list[NetworkInterfaceInfo]] = {
-        **state.node_networks,
-        event.node_id: event.network_interfaces,
-    }
-    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))
-    reconstructed = _reconstruct_profile(
-        event.node_id, state, network_interfaces=event.network_interfaces
-    )
-    topology.update_node_profile(event.node_id, reconstructed)
-    return state.model_copy(
-        update={
-            "node_networks": new_networks,
+            "node_profiles": new_profiles,
             "topology": topology,
             "last_seen": last_seen,
         }
@@ -343,26 +232,57 @@ def apply_node_network_measured(event: NodeNetworkMeasured, state: State) -> Sta
 
 
 def apply_node_memory_measured(event: NodeMemoryMeasured, state: State) -> State:
+    existing = state.node_profiles.get(event.node_id)
     topology = copy.copy(state.topology)
 
-    new_memories: Mapping[NodeId, MemoryPerformanceProfile] = {
-        **state.node_memories,
-        event.node_id: event.memory,
-    }
-    last_seen: Mapping[NodeId, datetime] = {
-        **state.last_seen,
-        event.node_id: datetime.fromisoformat(event.when),
+    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))
-    reconstructed = _reconstruct_profile(event.node_id, state, memory=event.memory)
-    topology.update_node_profile(event.node_id, reconstructed)
+    topology.update_node_profile(event.node_id, updated)
     return state.model_copy(
-        update={
-            "node_memories": new_memories,
-            "topology": topology,
-            "last_seen": last_seen,
-        }
+        update={"node_profiles": updated_profiles, "topology": topology}
     )
 
 

src/exo/shared/types/api.py

@@ -11,21 +11,10 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 
 FinishReason = Literal[
-    "stop", "length", "tool_calls", "content_filter", "function_call", "error"
+    "stop", "length", "tool_calls", "content_filter", "function_call"
 ]
 
 
-class ErrorInfo(BaseModel):
-    message: str
-    type: str
-    param: str | None = None
-    code: int
-
-
-class ErrorResponse(BaseModel):
-    error: ErrorInfo
-
-
 class ModelListModel(BaseModel):
     id: str
     object: str = "model"

src/exo/shared/types/chunks.py

@@ -22,7 +22,6 @@ class TokenChunk(BaseChunk):
     token_id: int
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
-    error_message: str | None = None
 
 
 class ImageChunk(BaseChunk):

src/exo/shared/types/events.py

@@ -2,14 +2,10 @@ from datetime import datetime
 
 from pydantic import Field
 
-from exo.shared.topology import Connection
+from exo.shared.topology import Connection, NodePerformanceProfile
 from exo.shared.types.chunks import GenerationChunk
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NetworkInterfaceInfo,
-    SystemPerformanceProfile,
-)
+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
@@ -89,35 +85,13 @@ class NodeTimedOut(BaseEvent):
     node_id: NodeId
 
 
-class NodeIdentityMeasured(BaseEvent):
-    """Static identity info - emitted once at startup."""
-
+class NodePerformanceMeasured(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
-    model_id: str
-    chip_id: str
-    friendly_name: str
-
-
-class NodeSystemMeasured(BaseEvent):
-    """Dynamic system metrics (GPU, temp, power) - emitted at 1s intervals."""
-
-    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
-    system: SystemPerformanceProfile
-
-
-class NodeNetworkMeasured(BaseEvent):
-    """Semi-static network interface info - emitted at 30s intervals."""
-
-    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
-    network_interfaces: list[NetworkInterfaceInfo]
+    node_profile: NodePerformanceProfile
 
 
 class NodeMemoryMeasured(BaseEvent):
-    """Dynamic memory metrics - emitted at 0.5s intervals."""
-
     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
@@ -153,9 +127,7 @@ Event = (
     | RunnerDeleted
     | NodeCreated
     | NodeTimedOut
-    | NodeIdentityMeasured
-    | NodeSystemMeasured
-    | NodeNetworkMeasured
+    | NodePerformanceMeasured
     | NodeMemoryMeasured
     | NodeDownloadProgress
     | ChunkGenerated

src/exo/shared/types/profiling.py

@@ -52,21 +52,13 @@ class NetworkInterfaceInfo(CamelCaseModel):
     ip_address: str
 
 
-class NodeIdentity(CamelCaseModel):
-    """Static identity info for a node."""
-
+class NodePerformanceProfile(CamelCaseModel):
     model_id: str
     chip_id: str
     friendly_name: str
-
-
-class NodePerformanceProfile(CamelCaseModel):
-    model_id: str | None = None
-    chip_id: str | None = None
-    friendly_name: str | None = None
-    memory: MemoryPerformanceProfile | None = None
+    memory: MemoryPerformanceProfile
     network_interfaces: list[NetworkInterfaceInfo] = []
-    system: SystemPerformanceProfile | None = None
+    system: SystemPerformanceProfile
 
 
 class ConnectionProfile(CamelCaseModel):

src/exo/shared/types/state.py

@@ -7,12 +7,7 @@ from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
 from exo.shared.types.common import NodeId
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NetworkInterfaceInfo,
-    NodeIdentity,
-    SystemPerformanceProfile,
-)
+from exo.shared.types.profiling import NodePerformanceProfile
 from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -40,10 +35,7 @@ class State(CamelCaseModel):
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
-    node_identities: Mapping[NodeId, NodeIdentity] = {}
-    node_memories: Mapping[NodeId, MemoryPerformanceProfile] = {}
-    node_systems: Mapping[NodeId, SystemPerformanceProfile] = {}
-    node_networks: Mapping[NodeId, list[NetworkInterfaceInfo]] = {}
+    node_profiles: Mapping[NodeId, NodePerformanceProfile] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)

src/exo/worker/download/download_utils.py

@@ -245,15 +245,12 @@ def create_http_session(
         sock_read_timeout = 1800
         sock_connect_timeout = 60
 
-    ssl_context = ssl.create_default_context(
-        cafile=os.getenv("SSL_CERT_FILE") or certifi.where()
-    )
+    ssl_context = ssl.create_default_context(cafile=certifi.where())
     connector = aiohttp.TCPConnector(ssl=ssl_context)
 
     return aiohttp.ClientSession(
         auto_decompress=auto_decompress,
         connector=connector,
-        proxy=os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY") or None,
         timeout=aiohttp.ClientTimeout(
             total=total_timeout,
             connect=connect_timeout,

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

@@ -46,11 +46,9 @@ class CustomMlxLayer(nn.Module):
 
     def __init__(self, original_layer: _LayerCallable):
         super().__init__()
+        # Set twice to avoid __setattr__ recursion
         object.__setattr__(self, "_original_layer", original_layer)
-
-    @property
-    def original_layer(self) -> _LayerCallable:
-        return cast(_LayerCallable, object.__getattribute__(self, "_original_layer"))
+        self.original_layer: _LayerCallable = original_layer
 
     # Calls __getattr__ for any attributes not found on nn.Module (e.g. use_sliding)
     if not TYPE_CHECKING:
@@ -60,7 +58,7 @@ class CustomMlxLayer(nn.Module):
                 return super().__getattr__(name)
             except AttributeError:
                 original_layer = object.__getattribute__(self, "_original_layer")
-                return getattr(original_layer, name)
+                return object.__getattribute__(original_layer, name)
 
 
 class PipelineFirstLayer(CustomMlxLayer):
@@ -170,21 +168,11 @@ def pipeline_auto_parallel(
         inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
             start_layer:end_layer
         ]
-        # We can assume the model has at least one layer thanks to placement.
-        # If a layer type doesn't exist, we can set it to 0.
-        inner_model_instance.swa_idx = (
-            0
-            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "sliding_attention"
-            )
+        inner_model_instance.swa_idx = inner_model_instance.layer_types.index(  # type: ignore
+            "sliding_attention"
         )
-        inner_model_instance.ga_idx = (
-            0
-            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "full_attention"
-            )
+        inner_model_instance.ga_idx = inner_model_instance.layer_types.index(  # type: ignore
+            "full_attention"
         )
 
     _set_layers(model, layers)

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

@@ -2,9 +2,7 @@ import json
 import os
 import resource
 import sys
-import threading
 import time
-from collections.abc import Callable
 from pathlib import Path
 from typing import Any, cast
 
@@ -84,45 +82,6 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     )
 
 
-class ModelLoadingTimeoutError(Exception):
-    pass
-
-
-TimeoutCallback = Callable[[], None]
-
-
-def eval_with_timeout(
-    mlx_item: Any,  # pyright: ignore[reportAny]
-    timeout_seconds: float = 60.0,
-    on_timeout: TimeoutCallback | None = None,
-) -> None:
-    """Evaluate MLX item with a hard timeout.
-
-    If on_timeout callback is provided, it will be called before terminating
-    the process. This allows the runner to send a failure event before exit.
-    """
-    completed = threading.Event()
-
-    def watchdog() -> None:
-        if not completed.wait(timeout=timeout_seconds):
-            logger.error(
-                f"mlx_item evaluation timed out after {timeout_seconds:.0f}s. "
-                "This may indicate an issue with FAST_SYNCH and tensor parallel sharding. "
-                "Terminating process."
-            )
-            if on_timeout is not None:
-                on_timeout()
-            os._exit(1)
-
-    watchdog_thread = threading.Thread(target=watchdog, daemon=True)
-    watchdog_thread.start()
-
-    try:
-        mx.eval(mlx_item)  # pyright: ignore[reportAny]
-    finally:
-        completed.set()
-
-
 def mx_barrier(group: Group | None = None):
     mx.eval(
         mx.distributed.all_sum(
@@ -229,9 +188,7 @@ def initialize_mlx(
 
 
 def load_mlx_items(
-    bound_instance: BoundInstance,
-    group: Group | None,
-    on_timeout: TimeoutCallback | None = None,
+    bound_instance: BoundInstance, group: Group | None
 ) -> tuple[Model, TokenizerWrapper]:
     if group is None:
         logger.info(f"Single device used for {bound_instance.instance}")
@@ -245,9 +202,7 @@ def load_mlx_items(
     else:
         logger.info("Starting distributed init")
         start_time = time.perf_counter()
-        model, tokenizer = shard_and_load(
-            bound_instance.bound_shard, group=group, on_timeout=on_timeout
-        )
+        model, tokenizer = shard_and_load(bound_instance.bound_shard, group=group)
         end_time = time.perf_counter()
         logger.info(
             f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
@@ -261,7 +216,6 @@ def load_mlx_items(
 def shard_and_load(
     shard_metadata: ShardMetadata,
     group: Group,
-    on_timeout: TimeoutCallback | None = None,
 ) -> tuple[nn.Module, TokenizerWrapper]:
     model_path = build_model_path(shard_metadata.model_meta.model_id)
 
@@ -298,15 +252,7 @@ def shard_and_load(
             logger.info(f"loading model from {model_path} with pipeline parallelism")
             model = pipeline_auto_parallel(model, group, shard_metadata)
 
-    # Estimate timeout based on model size
-    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "60"))
-    model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3)
-    timeout_seconds = base_timeout + model_size_gb / 5
-    logger.info(
-        f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
-        f"(model size: {model_size_gb:.1f}GB)"
-    )
-    eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
+    mx.eval(model.parameters())
 
     # TODO: Do we need this?
     mx.eval(model)

src/exo/worker/main.py

@@ -16,10 +16,8 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     NodeDownloadProgress,
-    NodeIdentityMeasured,
     NodeMemoryMeasured,
-    NodeNetworkMeasured,
-    NodeSystemMeasured,
+    NodePerformanceMeasured,
     TaskCreated,
     TaskStatusUpdated,
     TopologyEdgeCreated,
@@ -27,11 +25,7 @@ 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,
-    NetworkInterfaceInfo,
-    SystemPerformanceProfile,
-)
+from exo.shared.types.profiling import MemoryPerformanceProfile, NodePerformanceProfile
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     CreateRunner,
@@ -57,13 +51,7 @@ from exo.worker.download.download_utils import (
 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 (
-    IdentityMetrics,
-    start_polling_identity_metrics,
-    start_polling_memory_metrics,
-    start_polling_network_metrics,
-    start_polling_system_metrics,
-)
+from exo.worker.utils import start_polling_memory_metrics, start_polling_node_metrics
 from exo.worker.utils.net_profile import check_reachable
 
 
@@ -110,51 +98,37 @@ class Worker:
     async def run(self):
         logger.info("Starting Worker")
 
-        async def identity_callback(identity: IdentityMetrics) -> None:
+        # TODO: CLEANUP HEADER
+        async def resource_monitor_callback(
+            node_performance_profile: NodePerformanceProfile,
+        ) -> None:
             await self.event_sender.send(
-                NodeIdentityMeasured(
+                NodePerformanceMeasured(
                     node_id=self.node_id,
-                    model_id=identity.model_id,
-                    chip_id=identity.chip_id,
-                    friendly_name=identity.friendly_name,
+                    node_profile=node_performance_profile,
                     when=str(datetime.now(tz=timezone.utc)),
                 ),
             )
 
-        async def system_callback(system: SystemPerformanceProfile) -> None:
-            await self.event_sender.send(
-                NodeSystemMeasured(
-                    node_id=self.node_id,
-                    system=system,
-                    when=str(datetime.now(tz=timezone.utc)),
-                ),
-            )
-
-        async def network_callback(interfaces: list[NetworkInterfaceInfo]) -> None:
-            await self.event_sender.send(
-                NodeNetworkMeasured(
-                    node_id=self.node_id,
-                    network_interfaces=interfaces,
-                    when=str(datetime.now(tz=timezone.utc)),
-                ),
-            )
-
-        async def memory_callback(memory: MemoryPerformanceProfile) -> None:
+        async def memory_monitor_callback(
+            memory_profile: MemoryPerformanceProfile,
+        ) -> None:
             await self.event_sender.send(
                 NodeMemoryMeasured(
                     node_id=self.node_id,
-                    memory=memory,
+                    memory=memory_profile,
                     when=str(datetime.now(tz=timezone.utc)),
                 )
             )
 
+        # END CLEANUP
+
         async with create_task_group() as tg:
             self._tg = tg
             tg.start_soon(self.plan_step)
-            tg.start_soon(start_polling_identity_metrics, identity_callback)
-            tg.start_soon(start_polling_system_metrics, system_callback)
-            tg.start_soon(start_polling_network_metrics, network_callback)
-            tg.start_soon(start_polling_memory_metrics, memory_callback)
+            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)

src/exo/worker/runner/bootstrap.py

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

src/exo/worker/runner/runner.py

@@ -67,7 +67,6 @@ def main(
         bound_instance.bound_runner_id,
         bound_instance.bound_shard,
     )
-    device_rank = shard_metadata.device_rank
     logger.info("hello from the runner")
     if getattr(shard_metadata, "immediate_exception", False):
         raise Exception("Fake exception - runner failed to spin up.")
@@ -119,20 +118,7 @@ def main(
                         )
                     )
 
-                    def on_model_load_timeout() -> None:
-                        event_sender.send(
-                            RunnerStatusUpdated(
-                                runner_id=runner_id,
-                                runner_status=RunnerFailed(
-                                    error_message="Model loading timed out"
-                                ),
-                            )
-                        )
-                        time.sleep(0.5)
-
-                    model, tokenizer = load_mlx_items(
-                        bound_instance, group, on_timeout=on_model_load_timeout
-                    )
+                    model, tokenizer = load_mlx_items(bound_instance, group)
 
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
@@ -162,6 +148,8 @@ def main(
                 case ChatCompletion(task_params=task_params, command_id=command_id) if (
                     isinstance(current_status, RunnerReady)
                 ):
+                    assert model
+                    assert tokenizer
                     logger.info(f"received chat request: {str(task)[:500]}")
                     current_status = RunnerRunning()
                     logger.info("runner running")
@@ -170,61 +158,41 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
-                    assert model
-                    assert tokenizer
                     assert task_params.messages[0].content is not None
+                    _check_for_debug_prompts(task_params.messages[0].content)
 
-                    try:
-                        _check_for_debug_prompts(task_params.messages[0].content)
+                    # Generate responses using the actual MLX generation
+                    mlx_generator = mlx_generate(
+                        model=model,
+                        tokenizer=tokenizer,
+                        task=task_params,
+                    )
 
-                        # Generate responses using the actual MLX generation
-                        mlx_generator = mlx_generate(
-                            model=model,
-                            tokenizer=tokenizer,
-                            task=task_params,
-                        )
+                    # GPT-OSS specific parsing to match other model formats.
+                    if isinstance(model, GptOssModel):
+                        mlx_generator = parse_gpt_oss(mlx_generator)
 
-                        # GPT-OSS specific parsing to match other model formats.
-                        if isinstance(model, GptOssModel):
-                            mlx_generator = parse_gpt_oss(mlx_generator)
+                    # TODO: Add tool call parser here
 
-                        # TODO: Add tool call parser here
-
-                        for response in mlx_generator:
-                            match response:
-                                case GenerationResponse():
-                                    if device_rank == 0:
-                                        event_sender.send(
-                                            ChunkGenerated(
-                                                command_id=command_id,
-                                                chunk=TokenChunk(
-                                                    idx=response.token,
-                                                    model=shard_metadata.model_meta.model_id,
-                                                    text=response.text,
-                                                    token_id=response.token,
-                                                    finish_reason=response.finish_reason,
-                                                    stats=response.stats,
-                                                ),
-                                            )
+                    for response in mlx_generator:
+                        match response:
+                            case GenerationResponse():
+                                if shard_metadata.device_rank == 0:
+                                    event_sender.send(
+                                        ChunkGenerated(
+                                            command_id=command_id,
+                                            chunk=TokenChunk(
+                                                idx=response.token,
+                                                model=shard_metadata.model_meta.model_id,
+                                                text=response.text,
+                                                token_id=response.token,
+                                                finish_reason=response.finish_reason,
+                                                stats=response.stats,
+                                            ),
                                         )
-
-                    # can we make this more explicit?
-                    except Exception as e:
-                        if device_rank == 0:
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=TokenChunk(
-                                        idx=0,
-                                        model=shard_metadata.model_meta.model_id,
-                                        text="",
-                                        token_id=0,
-                                        finish_reason="error",
-                                        error_message=str(e),
-                                    ),
-                                )
-                            )
-                        raise
+                                    )
+                                # case TokenizedResponse():
+                                # TODO: something here ig
 
                     current_status = RunnerReady()
                     logger.info("runner ready")

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

@@ -1,202 +0,0 @@
-# type: ignore
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from exo.shared.constants import EXO_MODELS_DIR
-
-
-class MockLayer(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.custom_attr = "test_value"
-        self.use_sliding = True
-
-    def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
-        return x * 2
-
-
-@dataclass(frozen=True)
-class PipelineTestConfig:
-    model_path: Path
-    total_layers: int
-    base_port: int
-    max_tokens: int
-
-
-def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]:
-    import json
-    import tempfile
-
-    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
-
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
-        json.dump(hosts, f)
-        hostfile_path = f.name
-
-    return hostfile_path, hosts
-
-
-# Use GPT OSS 20b to test as it is a model with a lot of strange behaviour
-
-DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig(
-    model_path=EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8",
-    total_layers=24,
-    base_port=29600,
-    max_tokens=200,
-)
-
-
-def run_gpt_oss_pipeline_device(
-    rank: int,
-    world_size: int,
-    hostfile_path: str,
-    model_path: Path,
-    layer_splits: list[tuple[int, int]],
-    prompt_tokens: int,
-    prefill_step_size: int,
-    result_queue: Any,  # pyright: ignore[reportAny]
-    max_tokens: int = 200,
-) -> None:
-    import os
-    import traceback
-
-    os.environ["MLX_HOSTFILE"] = hostfile_path
-    os.environ["MLX_RANK"] = str(rank)
-
-    import mlx.core as mlx_core
-    from mlx_lm import load, stream_generate
-
-    from exo.shared.types.memory import Memory
-    from exo.shared.types.models import ModelId, ModelMetadata
-    from exo.shared.types.worker.shards import PipelineShardMetadata
-    from exo.worker.engines.mlx.auto_parallel import pipeline_auto_parallel
-
-    try:
-        group = mlx_core.distributed.init(backend="ring", strict=True)
-
-        model, tokenizer = load(str(model_path))
-
-        # Generate a prompt of exact token length
-        base_text = "The quick brown fox jumps over the lazy dog. "
-        base_tokens = tokenizer.encode(base_text)
-        base_len = len(base_tokens)
-
-        # Build prompt with approximate target length
-        repeats = (prompt_tokens // base_len) + 2
-        long_text = base_text * repeats
-        tokens = tokenizer.encode(long_text)
-        # Truncate to exact target length
-        tokens = tokens[:prompt_tokens]
-        prompt_text = tokenizer.decode(tokens)
-
-        formatted_prompt = tokenizer.apply_chat_template(
-            [{"role": "user", "content": prompt_text}],
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-        start_layer, end_layer = layer_splits[rank]
-
-        shard_meta = PipelineShardMetadata(
-            model_meta=ModelMetadata(
-                model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
-                pretty_name="GPT-OSS 20B",
-                storage_size=Memory.from_gb(12),
-                n_layers=24,
-                hidden_size=2880,
-                supports_tensor=False,
-            ),
-            device_rank=rank,
-            world_size=world_size,
-            start_layer=start_layer,
-            end_layer=end_layer,
-            n_layers=24,
-        )
-
-        model = pipeline_auto_parallel(model, group, shard_meta)
-
-        # Barrier before generation
-        barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
-        mlx_core.eval(barrier)
-
-        generated_text = ""
-        for response in stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=formatted_prompt,
-            max_tokens=max_tokens,
-            prefill_step_size=prefill_step_size,
-        ):
-            generated_text += response.text
-
-        result_queue.put((rank, True, generated_text))  # pyright: ignore[reportAny]
-
-    except Exception as e:
-        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))  # pyright: ignore[reportAny]
-
-
-def run_gpt_oss_tensor_parallel_device(
-    rank: int,
-    world_size: int,
-    hostfile_path: str,
-    model_path: Path,
-    prompt_tokens: int,
-    prefill_step_size: int,
-    result_queue: Any,  # pyright: ignore[reportAny]
-    max_tokens: int = 10,
-) -> None:
-    import os
-    import traceback
-
-    os.environ["MLX_HOSTFILE"] = hostfile_path
-    os.environ["MLX_RANK"] = str(rank)
-
-    import mlx.core as mlx_core
-    from mlx_lm import load, stream_generate
-
-    from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
-
-    try:
-        group = mlx_core.distributed.init(backend="ring", strict=True)
-
-        model, tokenizer = load(str(model_path))
-
-        base_text = "The quick brown fox jumps over the lazy dog. "
-        base_tokens = tokenizer.encode(base_text)
-        base_len = len(base_tokens)
-
-        repeats = (prompt_tokens // base_len) + 2
-        long_text = base_text * repeats
-        tokens = tokenizer.encode(long_text)
-        tokens = tokens[:prompt_tokens]
-        prompt_text = tokenizer.decode(tokens)
-
-        formatted_prompt = tokenizer.apply_chat_template(
-            [{"role": "user", "content": prompt_text}],
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-        model = tensor_auto_parallel(model, group)
-
-        barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
-        mlx_core.eval(barrier)
-
-        generated_text = ""
-        for response in stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=formatted_prompt,
-            max_tokens=max_tokens,
-            prefill_step_size=prefill_step_size,
-        ):
-            generated_text += response.text
-
-        result_queue.put((rank, True, generated_text))  # pyright: ignore[reportAny]
-
-    except Exception as e:
-        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))  # pyright: ignore[reportAny]

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

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

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

@@ -121,21 +121,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
 
 
-# Use a fake event_sender to remove test flakiness.
-class EventCollector:
-    def __init__(self) -> None:
-        self.events: list[Event] = []
-
-    def send(self, event: Event) -> None:
-        self.events.append(event)
-
-    def close(self) -> None:
-        pass
-
-    def join(self) -> None:
-        pass
-
-
 def _run(tasks: Iterable[Task]):
     bound_instance = get_bound_mlx_ring_instance(
         instance_id=INSTANCE_1_ID,
@@ -145,20 +130,22 @@ def _run(tasks: Iterable[Task]):
     )
 
     task_sender, task_receiver = mp_channel[Task]()
-    event_sender = EventCollector()
+    event_sender, event_receiver = mp_channel[Event]()
 
-    with task_sender:
+    with task_sender, event_receiver:
         for t in tasks:
             task_sender.send(t)
 
         # worst monkeypatch known to man
         # this is some c++ nonsense
+        event_sender.close = nothin
+        event_sender.join = nothin
         task_receiver.close = nothin
         task_receiver.join = nothin
 
-        mlx_runner.main(bound_instance, event_sender, task_receiver)  # type: ignore[arg-type]
+        mlx_runner.main(bound_instance, event_sender, task_receiver)
 
-        return event_sender.events
+        return event_receiver.collect()
 
 
 def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):

src/exo/worker/utils/__init__.py

@@ -1,15 +1,6 @@
-from .profile import (
-    IdentityMetrics,
-    start_polling_identity_metrics,
-    start_polling_memory_metrics,
-    start_polling_network_metrics,
-    start_polling_system_metrics,
-)
+from .profile import start_polling_memory_metrics, start_polling_node_metrics
 
 __all__ = [
-    "IdentityMetrics",
-    "start_polling_identity_metrics",
+    "start_polling_node_metrics",
     "start_polling_memory_metrics",
-    "start_polling_network_metrics",
-    "start_polling_system_metrics",
 ]

src/exo/worker/utils/net_profile.py

@@ -22,6 +22,7 @@ async def check_reachability(
     url = f"http://{target_ip}:52415/node_id"
 
     remote_node_id = None
+
     last_error = None
 
     for _ in range(REACHABILITY_ATTEMPTS):
@@ -39,24 +40,24 @@ async def check_reachability(
             remote_node_id = NodeId(body)
             break
 
-        # expected failure cases
         except (
-            httpx.TimeoutException,
-            httpx.NetworkError,
-        ):
-            await anyio.sleep(1)
-
-        # other failures should be logged on last attempt
-        except httpx.HTTPError as e:
+            httpx.ConnectError,
+            httpx.ConnectTimeout,
+            httpx.ReadTimeout,
+            httpx.RemoteProtocolError,
+        ) as e:
             last_error = e
             await anyio.sleep(1)
 
-    if last_error is not None:
-        logger.warning(
-            f"connect error {type(last_error).__name__} from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down"
-        )
-
-    if remote_node_id is None:
+    else:
+        if last_error is not None:
+            logger.warning(
+                f"connect error {type(last_error).__name__} from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down"
+            )
+        else:
+            logger.warning(
+                f"malformed response from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down"
+            )
         return
 
     if remote_node_id != expected_node_id:

src/exo/worker/utils/profile.py

@@ -1,7 +1,6 @@
 import asyncio
 import os
 import platform
-from dataclasses import dataclass
 from typing import Any, Callable, Coroutine
 
 import anyio
@@ -10,7 +9,7 @@ from loguru import logger
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryPerformanceProfile,
-    NetworkInterfaceInfo,
+    NodePerformanceProfile,
     SystemPerformanceProfile,
 )
 
@@ -28,13 +27,6 @@ from .system_info import (
 )
 
 
-@dataclass(frozen=True)
-class IdentityMetrics:
-    model_id: str
-    chip_id: str
-    friendly_name: str
-
-
 async def get_metrics_async() -> Metrics | None:
     """Return detailed Metrics on macOS or a minimal fallback elsewhere."""
 
@@ -75,73 +67,48 @@ async def start_polling_memory_metrics(
             await anyio.sleep(poll_interval_s)
 
 
-async def start_polling_identity_metrics(
-    callback: Callable[[IdentityMetrics], Coroutine[Any, Any, None]],
-    *,
-    poll_interval_s: float = 30.0,
-) -> None:
-    """Continuously poll and emit identity metrics at 30s intervals."""
-    while True:
-        try:
-            model_id, chip_id = await get_model_and_chip()
-            friendly_name = await get_friendly_name()
-            await callback(
-                IdentityMetrics(
-                    model_id=model_id,
-                    chip_id=chip_id,
-                    friendly_name=friendly_name,
-                )
-            )
-        except Exception as e:
-            logger.opt(exception=e).error("Failed to emit identity metrics")
-        finally:
-            await anyio.sleep(poll_interval_s)
-
-
-async def start_polling_system_metrics(
-    callback: Callable[[SystemPerformanceProfile], Coroutine[Any, Any, None]],
-    *,
-    poll_interval_s: float = 1.0,
-) -> None:
-    """Continuously poll and emit system metrics (GPU, temp, power) at 1s intervals."""
+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(
-                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,
+                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(
-                "[system_monitor] Operation timed out after 30s, skipping this cycle."
+                "[resource_monitor] Operation timed out after 30s, skipping this cycle."
             )
         except MacMonError as e:
-            logger.opt(exception=e).error("System Monitor encountered error")
+            logger.opt(exception=e).error("Resource Monitor encountered error")
             return
         finally:
             await anyio.sleep(poll_interval_s)
-
-
-async def start_polling_network_metrics(
-    callback: Callable[[list[NetworkInterfaceInfo]], Coroutine[Any, Any, None]],
-    *,
-    poll_interval_s: float = 30.0,
-) -> None:
-    """Continuously poll and emit network interface info at 30s intervals."""
-    while True:
-        try:
-            network_interfaces = get_network_interfaces()
-            await callback(network_interfaces)
-        except Exception as e:
-            logger.opt(exception=e).error("Network Monitor encountered error")
-        finally:
-            await anyio.sleep(poll_interval_s)