style: format app.svelte.ts with nix fmt

- Strip tokens (logprobs data) from messages before saving to localStorage
  since they're large and not essential for persistence
- Add pruneOldConversations() to automatically remove oldest conversations
  when quota is exceeded
- This prevents QuotaExceededError from crashing the app

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

README.md

@@ -42,7 +42,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
   <summary>Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
   <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-1-qwen3-235b.jpeg" alt="Benchmark - Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
   <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
   </p>
 </details>
 
@@ -50,7 +50,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
   <summary>DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
   <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-2-deepseek-3.1-671b.jpeg" alt="Benchmark - DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
   <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
   </p>
 </details>
 
@@ -58,7 +58,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
   <summary>Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
   <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-3-kimi-k2-thinking.jpeg" alt="Benchmark - Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
   <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
   </p>
 </details>
 
@@ -163,24 +163,6 @@ This starts the exo dashboard and API at http://localhost:52415/
 
 **Important note for Linux users:** Currently, exo runs on CPU on Linux. GPU support for Linux platforms is under development. If you'd like to see support for your specific Linux hardware, please [search for existing feature requests](https://github.com/exo-explore/exo/issues) or create a new one.
 
-**Configuration Options:**
-
-- `--no-worker`: Run exo without the worker component. Useful for coordinator-only nodes that handle networking and orchestration but don't execute inference tasks. This is helpful for machines without sufficient GPU resources but with good network connectivity.
-
-  ```bash
-  uv run exo --no-worker
-  ```
-
-**File Locations (Linux):**
-
-exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux:
-
-- **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
-- **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
-- **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
-
-You can override these locations by setting the corresponding XDG environment variables.
-
 ### macOS App
 
 exo ships a macOS app that runs in the background on your Mac.
@@ -193,19 +175,6 @@ Download the latest build here: [EXO-latest.dmg](https://assets.exolabs.net/EXO-
 
 The app will ask for permission to modify system settings and install a new Network profile. Improvements to this are being worked on.
 
-**Custom Namespace for Cluster Isolation:**
-
-The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. This is configured through the `EXO_LIBP2P_NAMESPACE` setting:
-
-- **Use cases**:
-  - Running multiple separate exo clusters on the same network
-  - Isolating development/testing clusters from production clusters
-  - Preventing accidental cluster joining
-
-- **Configuration**: Access this setting in the app's Advanced settings (or set the `EXO_LIBP2P_NAMESPACE` environment variable when running from source)
-
-The namespace is logged on startup for debugging purposes.
-
 #### Uninstalling the macOS App
 
 The recommended way to uninstall is through the app itself: click the menu bar icon → Advanced → Uninstall. This cleanly removes all system components.
@@ -352,52 +321,6 @@ For further details, see:
 
 ---
 
-## Benchmarking
-
-The `exo-bench` tool measures model prefill and token generation speed across different placement configurations. This helps you optimize model performance and validate improvements.
-
-**Prerequisites:**
-- Nodes should be running with `uv run exo` before benchmarking
-- The tool uses the `/bench/chat/completions` endpoint
-
-**Basic usage:**
-
-```bash
-uv run bench/exo_bench.py \
-  --model llama-3.2-1b \
-  --pp 128,256,512 \
-  --tg 128,256
-```
-
-**Key parameters:**
-
-- `--model`: Model to benchmark (short ID or HuggingFace ID)
-- `--pp`: Prompt size hints (comma-separated integers)
-- `--tg`: Generation lengths (comma-separated integers)
-- `--max-nodes`: Limit placements to N nodes (default: 4)
-- `--instance-meta`: Filter by `ring`, `jaccl`, or `both` (default: both)
-- `--sharding`: Filter by `pipeline`, `tensor`, or `both` (default: both)
-- `--repeat`: Number of repetitions per configuration (default: 1)
-- `--warmup`: Warmup runs per placement (default: 0)
-- `--json-out`: Output file for results (default: bench/results.json)
-
-**Example with filters:**
-
-```bash
-uv run bench/exo_bench.py \
-  --model llama-3.2-1b \
-  --pp 128,512 \
-  --tg 128 \
-  --max-nodes 2 \
-  --sharding tensor \
-  --repeat 3 \
-  --json-out my-results.json
-```
-
-The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
-
----
-
 ## Hardware Accelerator Support
 
 On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working on extending hardware accelerator support. If you'd like support for a new hardware platform, please [search for an existing feature request](https://github.com/exo-explore/exo/issues) and add a thumbs up so we know what hardware is important to the community.
@@ -406,4 +329,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working
 
 ## Contributing
 
-See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.

TODO.md

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

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -6,7 +6,7 @@ enum NetworkSetupHelper {
     private static let logger = Logger(subsystem: "io.exo.EXO", category: "NetworkSetup")
     private static let daemonLabel = "io.exo.networksetup"
     private static let scriptDestination =
-        "/Library/Application Support/EXO/disable_bridge.sh"
+        "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
     private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist"
     private static let requiredStartInterval: Int = 1791
 
@@ -28,6 +28,35 @@ enum NetworkSetupHelper {
         # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist
         /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true
 
+        networksetup -listlocations | grep -q exo || {
+          networksetup -createlocation exo
+        }
+
+        networksetup -switchtolocation exo
+        networksetup -listallhardwareports \\
+          | awk -F': ' '/Hardware Port: / {print $2}' \\
+          | while IFS=":" read -r name; do
+              case "$name" in
+                "Ethernet Adapter"*)
+                        ;;
+                "Thunderbolt Bridge")
+                        ;;
+                "Thunderbolt "*)
+                  networksetup -listallnetworkservices \\
+                    | grep -q "EXO $name" \\
+                      || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \\
+                      || continue
+                  networksetup -setdhcp "EXO $name"
+                        ;;
+                *)
+                  networksetup -listallnetworkservices \\
+                    | grep -q "$name" \\
+                      || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \\
+                      || continue
+                        ;;
+              esac
+            done
+
         networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
           networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
         } || true
@@ -112,13 +141,6 @@ enum NetworkSetupHelper {
         let scriptExists = manager.fileExists(atPath: scriptDestination)
         let plistExists = manager.fileExists(atPath: plistDestination)
         guard scriptExists, plistExists else { return false }
-        guard
-            let installedScript = try? String(contentsOfFile: scriptDestination, encoding: .utf8),
-            installedScript.trimmingCharacters(in: .whitespacesAndNewlines)
-                == setupScript.trimmingCharacters(in: .whitespacesAndNewlines)
-        else {
-            return false
-        }
         guard
             let data = try? Data(contentsOf: URL(fileURLWithPath: plistDestination)),
             let plist = try? PropertyListSerialization.propertyList(

bench/exo_bench.py

@@ -16,6 +16,9 @@ from urllib.parse import urlencode
 from loguru import logger
 from transformers import AutoTokenizer
 
+from exo.shared.models.model_cards import MODEL_CARDS
+from exo.shared.types.memory import Memory
+
 
 class ExoHttpError(RuntimeError):
     def __init__(self, status: int, reason: str, body_preview: str):
@@ -487,17 +490,17 @@ def main() -> int:
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
             for pp in pp_list:
-                # if (
-                #     pp * n_nodes > 2048
-                #     and "ring" in instance_meta.lower()
-                #     and "tensor" in sharding.lower()
-                # ):
-                #     model_card = MODEL_CARDS[short_id]
-                #     if model_card.metadata.storage_size > Memory.from_gb(10):
-                #         logger.info(
-                #             f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
-                #         )
-                #         continue
+                if (
+                    pp * n_nodes > 2048
+                    and "ring" in instance_meta.lower()
+                    and "tensor" in sharding.lower()
+                ):
+                    model_card = MODEL_CARDS[short_id]
+                    if model_card.metadata.storage_size > Memory.from_gb(10):
+                        logger.info(
+                            f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
+                        )
+                        continue
                 for tg in tg_list:
                     runs: list[dict[str, Any]] = []
                     for r in range(args.repeat):

dashboard/package-lock.json

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

dashboard/src/lib/components/ChatMessages.svelte

@@ -1,14 +1,16 @@
 <script lang="ts">
-	import { 
-		messages, 
-		currentResponse, 
+	import {
+		messages,
+		currentResponse,
 		isLoading,
 		deleteMessage,
 		editAndRegenerate,
-		regenerateLastResponse
+		regenerateLastResponse,
+		regenerateFromToken
 	} from '$lib/stores/app.svelte';
 	import type { MessageAttachment } from '$lib/stores/app.svelte';
 	import MarkdownContent from './MarkdownContent.svelte';
+	import TokenHeatmap from './TokenHeatmap.svelte';
 
 	interface Props {
 		class?: string;
@@ -95,6 +97,23 @@
 let copiedMessageId = $state<string | null>(null);
 let expandedThinkingMessageIds = $state<Set<string>>(new Set());
 
+// Uncertainty view state - tracks which messages show token heatmap
+let uncertaintyViewMessageIds = $state<Set<string>>(new Set());
+
+function toggleUncertaintyView(messageId: string) {
+	const newSet = new Set(uncertaintyViewMessageIds);
+	if (newSet.has(messageId)) {
+		newSet.delete(messageId);
+	} else {
+		newSet.add(messageId);
+	}
+	uncertaintyViewMessageIds = newSet;
+}
+
+function isUncertaintyViewEnabled(messageId: string): boolean {
+	return uncertaintyViewMessageIds.has(messageId);
+}
+
 	function formatTimestamp(timestamp: number): string {
 		return new Date(timestamp).toLocaleTimeString('en-US', { 
 			hour12: false,
@@ -366,7 +385,17 @@ function isThinkingExpanded(messageId: string): boolean {
 									</div>
 								{/if}
 								<div class="text-xs text-foreground">
-									<MarkdownContent content={message.content || (loading ? response : '')} />
+									{#if message.role === 'assistant' && isUncertaintyViewEnabled(message.id) && message.tokens && message.tokens.length > 0}
+										<!-- Uncertainty heatmap view -->
+										<TokenHeatmap
+											tokens={message.tokens}
+											isGenerating={loading}
+											onRegenerateFrom={(tokenIndex) => regenerateFromToken(message.id, tokenIndex)}
+										/>
+									{:else}
+										<!-- Normal markdown view -->
+										<MarkdownContent content={message.content || (loading ? response : '')} />
+									{/if}
 									{#if loading && !message.content}
 										<span class="inline-block w-2 h-4 bg-exo-yellow/70 ml-1 cursor-blink"></span>
 									{/if}
@@ -419,7 +448,20 @@ function isThinkingExpanded(messageId: string): boolean {
 								</svg>
 							</button>
 						{/if}
-						
+
+						<!-- Uncertainty view toggle (assistant messages with tokens only) -->
+						{#if message.role === 'assistant' && message.tokens && message.tokens.length > 0}
+							<button
+								onclick={() => toggleUncertaintyView(message.id)}
+								class="p-1.5 transition-colors rounded cursor-pointer {isUncertaintyViewEnabled(message.id) ? 'text-exo-yellow' : 'text-exo-light-gray hover:text-exo-yellow'}"
+								title={isUncertaintyViewEnabled(message.id) ? 'Hide uncertainty' : 'Show uncertainty'}
+							>
+								<svg class="w-3.5 h-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+									<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 19v-6a2 2 0 00-2-2H5a2 2 0 00-2 2v6a2 2 0 002 2h2a2 2 0 002-2zm0 0V9a2 2 0 012-2h2a2 2 0 012 2v10m-6 0a2 2 0 002 2h2a2 2 0 002-2m0 0V5a2 2 0 012-2h2a2 2 0 012 2v14a2 2 0 01-2 2h-2a2 2 0 01-2-2z" />
+								</svg>
+							</button>
+						{/if}
+
 						<!-- Delete button -->
 						<button
 							onclick={() => handleDeleteClick(message.id)}

dashboard/src/lib/components/MarkdownContent.svelte

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

dashboard/src/lib/components/ModelCard.svelte

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

dashboard/src/lib/components/TokenHeatmap.svelte

@@ -0,0 +1,192 @@
+<script lang="ts">
+	import type { TokenData } from '$lib/stores/app.svelte';
+
+	interface Props {
+		tokens: TokenData[];
+		class?: string;
+		isGenerating?: boolean;
+		onRegenerateFrom?: (tokenIndex: number) => void;
+	}
+
+	let { tokens, class: className = '', isGenerating = false, onRegenerateFrom }: Props = $props();
+
+	// Tooltip state - track both token data and index
+	let hoveredTokenIndex = $state<number | null>(null);
+	let hoveredPosition = $state<{ x: number; y: number } | null>(null);
+	let isTooltipHovered = $state(false);
+	let hideTimeoutId: ReturnType<typeof setTimeout> | null = null;
+
+	// Derive the hovered token from the index (stable across re-renders)
+	const hoveredToken = $derived(
+		hoveredTokenIndex !== null && hoveredPosition && tokens[hoveredTokenIndex]
+			? { token: tokens[hoveredTokenIndex], index: hoveredTokenIndex, ...hoveredPosition }
+			: null
+	);
+
+	/**
+	 * Get confidence styling based on probability.
+	 * Following Apple design principles: high confidence tokens blend in,
+	 * only uncertainty draws attention.
+	 */
+	function getConfidenceClass(probability: number): string {
+		if (probability > 0.8) return 'text-inherit'; // Expected tokens - blend in
+		if (probability > 0.5) return 'bg-gray-500/10 text-inherit'; // Slight hint
+		if (probability > 0.2) return 'bg-amber-500/15 text-amber-200/90'; // Subtle warmth
+		return 'bg-red-500/20 text-red-200/90'; // Draws attention
+	}
+
+	/**
+	 * Get border/underline styling for uncertain tokens
+	 */
+	function getBorderClass(probability: number): string {
+		if (probability > 0.8) return 'border-transparent'; // No border for expected
+		if (probability > 0.5) return 'border-gray-500/20';
+		if (probability > 0.2) return 'border-amber-500/30';
+		return 'border-red-500/40';
+	}
+
+	function clearHideTimeout() {
+		if (hideTimeoutId) {
+			clearTimeout(hideTimeoutId);
+			hideTimeoutId = null;
+		}
+	}
+
+	function handleMouseEnter(event: MouseEvent, token: TokenData, index: number) {
+		clearHideTimeout();
+		const rect = (event.target as HTMLElement).getBoundingClientRect();
+		hoveredTokenIndex = index;
+		hoveredPosition = {
+			x: rect.left + rect.width / 2,
+			y: rect.top - 10
+		};
+	}
+
+	function handleMouseLeave() {
+		clearHideTimeout();
+		// Use longer delay during generation to account for re-renders
+		const delay = isGenerating ? 300 : 100;
+		hideTimeoutId = setTimeout(() => {
+			if (!isTooltipHovered) {
+				hoveredTokenIndex = null;
+				hoveredPosition = null;
+			}
+		}, delay);
+	}
+
+	function handleTooltipEnter() {
+		clearHideTimeout();
+		isTooltipHovered = true;
+	}
+
+	function handleTooltipLeave() {
+		isTooltipHovered = false;
+		hoveredTokenIndex = null;
+		hoveredPosition = null;
+	}
+
+	function handleRegenerate() {
+		if (hoveredToken && onRegenerateFrom) {
+			const indexToRegenerate = hoveredToken.index;
+			// Clear hover state immediately
+			hoveredTokenIndex = null;
+			hoveredPosition = null;
+			isTooltipHovered = false;
+			// Call regenerate
+			onRegenerateFrom(indexToRegenerate);
+		}
+	}
+
+	function formatProbability(prob: number): string {
+		return (prob * 100).toFixed(1) + '%';
+	}
+
+	function formatLogprob(logprob: number): string {
+		return logprob.toFixed(3);
+	}
+
+	function getProbabilityColor(probability: number): string {
+		if (probability > 0.8) return 'text-gray-300';
+		if (probability > 0.5) return 'text-gray-400';
+		if (probability > 0.2) return 'text-amber-400';
+		return 'text-red-400';
+	}
+</script>
+
+<div class="token-heatmap leading-relaxed {className}">
+	{#each tokens as tokenData, i (i)}
+		<span
+			role="button"
+			tabindex="0"
+			class="token-span inline rounded px-0.5 py-0.5 cursor-pointer transition-all duration-150 border {getConfidenceClass(tokenData.probability)} {getBorderClass(tokenData.probability)} hover:opacity-80"
+			onmouseenter={(e) => handleMouseEnter(e, tokenData, i)}
+			onmouseleave={handleMouseLeave}
+		>{tokenData.token}</span>
+	{/each}
+</div>
+
+<!-- Tooltip -->
+{#if hoveredToken}
+	<div
+		class="fixed z-50"
+		style="left: {hoveredToken.x}px; top: {hoveredToken.y}px; transform: translate(-50%, -100%);"
+		onmouseenter={handleTooltipEnter}
+		onmouseleave={handleTooltipLeave}
+	>
+		<div class="bg-gray-900/95 backdrop-blur-sm border border-gray-700/50 rounded-xl shadow-xl p-3 text-sm min-w-48">
+			<!-- Token info -->
+			<div class="mb-2">
+				<span class="text-gray-500 text-xs">Token:</span>
+				<span class="text-white font-mono ml-1">"{hoveredToken.token.token}"</span>
+				<span class="{getProbabilityColor(hoveredToken.token.probability)} ml-2">{formatProbability(hoveredToken.token.probability)}</span>
+			</div>
+
+			<div class="text-gray-400 text-xs mb-1">
+				logprob: <span class="text-gray-300 font-mono">{formatLogprob(hoveredToken.token.logprob)}</span>
+			</div>
+
+			<!-- Top alternatives -->
+			{#if hoveredToken.token.topLogprobs.length > 0}
+				<div class="border-t border-gray-700/50 mt-2 pt-2">
+					<div class="text-gray-500 text-xs mb-1">Alternatives:</div>
+					{#each hoveredToken.token.topLogprobs.slice(0, 5) as alt, idx (idx)}
+						{@const altProb = Math.exp(alt.logprob)}
+						<div class="flex justify-between items-center text-xs py-0.5">
+							<span class="text-gray-300 font-mono truncate max-w-24">"{alt.token}"</span>
+							<span class="text-gray-400 ml-2">{formatProbability(altProb)}</span>
+						</div>
+					{/each}
+				</div>
+			{/if}
+
+			<!-- Regenerate button -->
+			{#if onRegenerateFrom}
+				<button
+					onclick={handleRegenerate}
+					class="w-full mt-2 pt-2 border-t border-gray-700/50 flex items-center justify-center gap-1.5 text-xs text-gray-400 hover:text-white transition-colors cursor-pointer"
+				>
+					<svg class="w-3 h-3" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+						<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
+					</svg>
+					Regenerate from here
+				</button>
+			{/if}
+		</div>
+		<!-- Arrow -->
+		<div class="absolute left-1/2 -translate-x-1/2 top-full">
+			<div class="border-8 border-transparent border-t-gray-900"></div>
+		</div>
+	</div>
+{/if}
+
+<style>
+	.token-heatmap {
+		word-wrap: break-word;
+		white-space: pre-wrap;
+	}
+
+	.token-span {
+		margin: 0;
+		border-width: 1px;
+	}
+</style>

dashboard/src/lib/components/TopologyGraph.svelte

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

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

@@ -71,66 +71,52 @@ export interface Instance {
 	};
 }
 
-// Granular node state types from the new state structure
-interface RawNodeIdentity {
+interface RawNodeProfile {
 	modelId?: string;
 	chipId?: string;
 	friendlyName?: string;
-}
-
-interface RawMemoryUsage {
-	ramTotal?: { inBytes: number };
-	ramAvailable?: { inBytes: number };
-	swapTotal?: { inBytes: number };
-	swapAvailable?: { inBytes: number };
-}
-
-interface RawSystemPerformanceProfile {
-	gpuUsage?: number;
-	temp?: number;
-	sysPower?: number;
-	pcpuUsage?: number;
-	ecpuUsage?: number;
-}
-
-interface RawNetworkInterfaceInfo {
-	name?: string;
-	ipAddress?: string;
-	addresses?: Array<{ address?: string } | string>;
-	ipv4?: string;
-	ipv6?: string;
-	ipAddresses?: string[];
-	ips?: string[];
-}
-
-interface RawNodeNetworkInfo {
-	interfaces?: RawNetworkInterfaceInfo[];
-}
-
-interface RawSocketConnection {
-	sinkMultiaddr?: {
-		address?: string;
-		ip_address?: string;
-		address_type?: string;
-		port?: number;
+	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 RawRDMAConnection {
-	sourceRdmaIface?: string;
-	sinkRdmaIface?: string;
+interface RawTopologyNode {
+	nodeId: string;
+	nodeProfile: RawNodeProfile;
 }
 
-type RawConnectionEdge = RawSocketConnection | RawRDMAConnection;
-
-// New nested mapping format: { source: { sink: [edge1, edge2, ...] } }
-type RawConnectionsMap = Record<string, Record<string, RawConnectionEdge[]>>;
+interface RawTopologyConnection {
+	localNodeId: string;
+	sendBackNodeId: string;
+	sendBackMultiaddr?:
+		| { multiaddr?: string; address?: string; ip_address?: string }
+		| string;
+}
 
 interface RawTopology {
-	nodes: string[];
-	connections?: RawConnectionsMap;
+	nodes: RawTopologyNode[];
+	connections?: RawTopologyConnection[];
 }
 
+type RawNodeProfiles = Record<string, RawNodeProfile>;
+
 export interface DownloadProgress {
 	totalBytes: number;
 	downloadedBytes: number;
@@ -185,11 +171,7 @@ interface RawStateResponse {
 	>;
 	runners?: Record<string, unknown>;
 	downloads?: Record<string, unknown[]>;
-	// New granular node state fields
-	nodeIdentities?: Record<string, RawNodeIdentity>;
-	nodeMemory?: Record<string, RawMemoryUsage>;
-	nodeSystem?: Record<string, RawSystemPerformanceProfile>;
-	nodeNetwork?: Record<string, RawNodeNetworkInfo>;
+	nodeProfiles?: RawNodeProfiles;
 }
 
 export interface MessageAttachment {
@@ -200,6 +182,20 @@ export interface MessageAttachment {
 	mimeType?: string;
 }
 
+// Token-level data for uncertainty visualization
+export interface TopLogprob {
+	token: string;
+	logprob: number;
+	bytes?: number[];
+}
+
+export interface TokenData {
+	token: string;
+	logprob: number;
+	probability: number; // exp(logprob)
+	topLogprobs: TopLogprob[];
+}
+
 export interface Message {
 	id: string;
 	role: "user" | "assistant" | "system";
@@ -209,6 +205,7 @@ export interface Message {
 	attachments?: MessageAttachment[];
 	ttftMs?: number; // Time to first token in ms (for assistant messages)
 	tps?: number; // Tokens per second (for assistant messages)
+	tokens?: TokenData[]; // Token-level data for uncertainty visualization
 }
 
 export interface Conversation {
@@ -224,69 +221,56 @@ export interface Conversation {
 
 const STORAGE_KEY = "exo-conversations";
 
-interface GranularNodeState {
-	nodeIdentities?: Record<string, RawNodeIdentity>;
-	nodeMemory?: Record<string, RawMemoryUsage>;
-	nodeSystem?: Record<string, RawSystemPerformanceProfile>;
-	nodeNetwork?: Record<string, RawNodeNetworkInfo>;
-}
-
-function transformNetworkInterface(iface: RawNetworkInterfaceInfo): {
-	name?: string;
-	addresses: string[];
-} {
-	const addresses: string[] = [];
-	if (iface.ipAddress && typeof iface.ipAddress === "string") {
-		addresses.push(iface.ipAddress);
-	}
-	if (Array.isArray(iface.addresses)) {
-		for (const addr of iface.addresses) {
-			if (typeof addr === "string") addresses.push(addr);
-			else if (addr && typeof addr === "object" && addr.address)
-				addresses.push(addr.address);
-		}
-	}
-	if (Array.isArray(iface.ipAddresses)) {
-		addresses.push(
-			...iface.ipAddresses.filter((a): a is string => typeof a === "string"),
-		);
-	}
-	if (Array.isArray(iface.ips)) {
-		addresses.push(
-			...iface.ips.filter((a): a is string => typeof a === "string"),
-		);
-	}
-	if (iface.ipv4 && typeof iface.ipv4 === "string") addresses.push(iface.ipv4);
-	if (iface.ipv6 && typeof iface.ipv6 === "string") addresses.push(iface.ipv6);
-
-	return {
-		name: iface.name,
-		addresses: Array.from(new Set(addresses)),
-	};
-}
-
 function transformTopology(
 	raw: RawTopology,
-	granularState: GranularNodeState,
+	profiles?: RawNodeProfiles,
 ): TopologyData {
 	const nodes: Record<string, NodeInfo> = {};
 	const edges: TopologyEdge[] = [];
 
-	for (const nodeId of raw.nodes || []) {
-		if (!nodeId) continue;
-
-		// Get data from granular state mappings
-		const identity = granularState.nodeIdentities?.[nodeId];
-		const memory = granularState.nodeMemory?.[nodeId];
-		const system = granularState.nodeSystem?.[nodeId];
-		const network = granularState.nodeNetwork?.[nodeId];
-
-		const ramTotal = memory?.ramTotal?.inBytes ?? 0;
-		const ramAvailable = memory?.ramAvailable?.inBytes ?? 0;
+	for (const node of raw.nodes || []) {
+		const mergedProfile = profiles?.[node.nodeId];
+		const profile = { ...(node.nodeProfile ?? {}), ...(mergedProfile ?? {}) };
+		const ramTotal = profile?.memory?.ramTotal?.inBytes ?? 0;
+		const ramAvailable = profile?.memory?.ramAvailable?.inBytes ?? 0;
 		const ramUsage = Math.max(ramTotal - ramAvailable, 0);
 
-		const rawInterfaces = network?.interfaces || [];
-		const networkInterfaces = rawInterfaces.map(transformNetworkInterface);
+		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) {
@@ -295,10 +279,10 @@ function transformTopology(
 			}
 		}
 
-		nodes[nodeId] = {
+		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,
@@ -309,42 +293,43 @@ 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,
 		};
 	}
 
-	// Handle connections - nested mapping format { source: { sink: [edges] } }
-	const connections = raw.connections;
-	if (connections && typeof connections === "object") {
-		for (const [source, sinks] of Object.entries(connections)) {
-			if (!sinks || typeof sinks !== "object") continue;
-			for (const [sink, edgeList] of Object.entries(sinks)) {
-				if (!Array.isArray(edgeList)) continue;
-				for (const edge of edgeList) {
-					let sendBackIp: string | undefined;
-					if (edge && typeof edge === "object" && "sinkMultiaddr" in edge) {
-						const multiaddr = edge.sinkMultiaddr;
-						if (multiaddr) {
-							sendBackIp =
-								multiaddr.ip_address ||
-								extractIpFromMultiaddr(multiaddr.address);
-						}
-					}
+	for (const conn of raw.connections || []) {
+		if (!conn.localNodeId || !conn.sendBackNodeId) continue;
+		if (conn.localNodeId === conn.sendBackNodeId) continue;
+		if (!nodes[conn.localNodeId] || !nodes[conn.sendBackNodeId]) continue;
 
-					if (nodes[source] && nodes[sink] && source !== sink) {
-						edges.push({ source, target: sink, sendBackIp });
-					}
-				}
+		let sendBackIp: string | undefined;
+		if (conn.sendBackMultiaddr) {
+			const multi = conn.sendBackMultiaddr;
+			if (typeof multi === "string") {
+				sendBackIp = extractIpFromMultiaddr(multi);
+			} else {
+				sendBackIp =
+					multi.ip_address ||
+					extractIpFromMultiaddr(multi.multiaddr) ||
+					extractIpFromMultiaddr(multi.address);
 			}
 		}
+
+		edges.push({
+			source: conn.localNodeId,
+			target: conn.sendBackNodeId,
+			sendBackIp,
+		});
 	}
 
 	return { nodes, edges };
@@ -398,6 +383,21 @@ class AppStore {
 	private fetchInterval: ReturnType<typeof setInterval> | null = null;
 	private previewsInterval: ReturnType<typeof setInterval> | null = null;
 	private lastConversationPersistTs = 0;
+	private currentRequestController: AbortController | null = null;
+
+	/**
+	 * Abort any in-flight generation request
+	 */
+	abortCurrentRequest(): boolean {
+		if (this.currentRequestController) {
+			this.currentRequestController.abort();
+			this.currentRequestController = null;
+			this.isLoading = false;
+			this.currentResponse = "";
+			return true;
+		}
+		return false;
+	}
 
 	constructor() {
 		if (browser) {
@@ -435,12 +435,61 @@ class AppStore {
 
 	/**
 	 * Save conversations to localStorage
+	 * Note: We strip tokens (logprobs data) to save space - they're large and not essential for persistence
 	 */
 	private saveConversationsToStorage() {
 		try {
-			localStorage.setItem(STORAGE_KEY, JSON.stringify(this.conversations));
+			// Strip tokens from messages to save localStorage space
+			const conversationsToSave = this.conversations.map((conv) => ({
+				...conv,
+				messages: conv.messages.map((msg) => {
+					// eslint-disable-next-line @typescript-eslint/no-unused-vars
+					const { tokens, ...msgWithoutTokens } = msg;
+					return msgWithoutTokens;
+				}),
+			}));
+			localStorage.setItem(STORAGE_KEY, JSON.stringify(conversationsToSave));
 		} catch (error) {
 			console.error("Failed to save conversations:", error);
+			// If quota exceeded, try to clear old conversations and retry
+			if (
+				error instanceof DOMException &&
+				error.name === "QuotaExceededError"
+			) {
+				console.warn(
+					"Storage quota exceeded, clearing oldest conversations...",
+				);
+				this.pruneOldConversations();
+			}
+		}
+	}
+
+	/**
+	 * Remove oldest conversations to free up storage space
+	 */
+	private pruneOldConversations() {
+		if (this.conversations.length <= 1) return;
+
+		// Sort by updatedAt and remove oldest half
+		const sorted = [...this.conversations].sort(
+			(a, b) => (b.updatedAt || 0) - (a.updatedAt || 0),
+		);
+		const keepCount = Math.max(1, Math.ceil(sorted.length / 2));
+		this.conversations = sorted.slice(0, keepCount);
+
+		// Try saving again
+		try {
+			const conversationsToSave = this.conversations.map((conv) => ({
+				...conv,
+				messages: conv.messages.map((msg) => {
+					// eslint-disable-next-line @typescript-eslint/no-unused-vars
+					const { tokens, ...msgWithoutTokens } = msg;
+					return msgWithoutTokens;
+				}),
+			}));
+			localStorage.setItem(STORAGE_KEY, JSON.stringify(conversationsToSave));
+		} catch {
+			console.error("Still failed to save after pruning");
 		}
 	}
 
@@ -898,12 +947,7 @@ class AppStore {
 			const data: RawStateResponse = await response.json();
 
 			if (data.topology) {
-				this.topologyData = transformTopology(data.topology, {
-					nodeIdentities: data.nodeIdentities,
-					nodeMemory: data.nodeMemory,
-					nodeSystem: data.nodeSystem,
-					nodeNetwork: data.nodeNetwork,
-				});
+				this.topologyData = transformTopology(data.topology, data.nodeProfiles);
 			}
 			if (data.instances) {
 				this.instances = data.instances;
@@ -1366,6 +1410,11 @@ class AppStore {
 		const assistantMessage = this.addMessage("assistant", "");
 		this.updateActiveConversation();
 
+		// Create abort controller for this request - must be defined before try block
+		// so it's available in the finally block
+		const controller = new AbortController();
+		this.currentRequestController = controller;
+
 		try {
 			// Build the messages array for the API with system prompt
 			const systemPrompt = {
@@ -1443,7 +1492,10 @@ class AppStore {
 					messages: apiMessages,
 					temperature: 0.7,
 					stream: true,
+					logprobs: true,
+					top_logprobs: 5,
 				}),
+				signal: controller.signal,
 			});
 
 			if (!response.ok) {
@@ -1459,6 +1511,7 @@ class AppStore {
 			const decoder = new TextDecoder();
 			let fullContent = "";
 			let buffer = "";
+			const collectedTokens: TokenData[] = [];
 
 			while (true) {
 				const { done, value } = await reader.read();
@@ -1480,8 +1533,8 @@ class AppStore {
 
 						try {
 							const parsed = JSON.parse(data);
-							const tokenContent = parsed.choices?.[0]?.delta?.content;
-							if (tokenContent) {
+							const delta = parsed.choices?.[0]?.delta?.content;
+							if (delta) {
 								// Track first token for TTFT
 								if (firstTokenTime === null) {
 									firstTokenTime = performance.now();
@@ -1498,7 +1551,30 @@ class AppStore {
 									this.tps = (tokenCount / elapsed) * 1000;
 								}
 
-								fullContent += tokenContent;
+								// Extract logprobs for uncertainty visualization
+								const logprobsData = parsed.choices?.[0]?.logprobs;
+								if (logprobsData?.content?.[0]) {
+									const logprobItem = logprobsData.content[0];
+									const tokenData: TokenData = {
+										token: logprobItem.token || delta,
+										logprob: logprobItem.logprob ?? 0,
+										probability: Math.exp(logprobItem.logprob ?? 0),
+										topLogprobs: (logprobItem.top_logprobs || []).map(
+											(item: {
+												token: string;
+												logprob: number;
+												bytes?: number[];
+											}) => ({
+												token: item.token,
+												logprob: item.logprob,
+												bytes: item.bytes,
+											}),
+										),
+									};
+									collectedTokens.push(tokenData);
+								}
+
+								fullContent += delta;
 
 								// Strip thinking tags for display and extract thinking content
 								const { displayContent, thinkingContent } =
@@ -1512,6 +1588,7 @@ class AppStore {
 								if (idx !== -1) {
 									this.messages[idx].content = displayContent;
 									this.messages[idx].thinking = thinkingContent || undefined;
+									this.messages[idx].tokens = [...collectedTokens];
 								}
 								this.persistActiveConversation();
 							}
@@ -1559,9 +1636,16 @@ class AppStore {
 				if (this.tps !== null) {
 					this.messages[idx].tps = this.tps;
 				}
+				if (collectedTokens.length > 0) {
+					this.messages[idx].tokens = collectedTokens;
+				}
 			}
 			this.persistActiveConversation();
 		} catch (error) {
+			// Don't show error for aborted requests (user cancelled)
+			if (error instanceof Error && error.name === "AbortError") {
+				return;
+			}
 			console.error("Error sending message:", error);
 			// Update the assistant message with error
 			const idx = this.messages.findIndex((m) => m.id === assistantMessage.id);
@@ -1571,6 +1655,237 @@ class AppStore {
 			}
 			this.persistActiveConversation();
 		} finally {
+			// Clean up controller if this is still the active request
+			if (this.currentRequestController === controller) {
+				this.currentRequestController = null;
+			}
+			this.isLoading = false;
+			this.currentResponse = "";
+			this.updateActiveConversation();
+		}
+	}
+
+	/**
+	 * Regenerate from a specific token in an assistant message.
+	 * Keeps content up to and including the specified token, then continues generation.
+	 * If a generation is already in progress, it will be aborted first.
+	 */
+	async regenerateFromToken(
+		messageId: string,
+		tokenIndex: number,
+	): Promise<void> {
+		// Abort any in-flight request first
+		this.abortCurrentRequest();
+
+		const messageIdx = this.messages.findIndex((m) => m.id === messageId);
+		if (messageIdx === -1) return;
+
+		const message = this.messages[messageIdx];
+		if (message.role !== "assistant" || !message.tokens) return;
+
+		// Get tokens up to and including the specified index
+		const tokensToKeep = message.tokens.slice(0, tokenIndex + 1);
+		const prefixText = tokensToKeep.map((t) => t.token).join("");
+
+		// Remove all messages after this assistant message
+		this.messages = this.messages.slice(0, messageIdx + 1);
+
+		// Update the message to show the prefix
+		this.messages[messageIdx].content = prefixText;
+		this.messages[messageIdx].tokens = tokensToKeep;
+
+		// Set up for continuation
+		this.isLoading = true;
+		this.currentResponse = prefixText;
+		this.ttftMs = null;
+		this.tps = null;
+		this.totalTokens = tokensToKeep.length;
+
+		// Create abort controller before try block so it's available in finally
+		const controller = new AbortController();
+		this.currentRequestController = controller;
+
+		try {
+			// Build messages for API - include the partial assistant message
+			const systemPrompt = {
+				role: "system" as const,
+				content:
+					"You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process.",
+			};
+
+			// Get all messages up to and including the one we're regenerating from
+			const apiMessages = [
+				systemPrompt,
+				...this.messages.map((m) => {
+					let msgContent = m.content;
+					if (m.attachments) {
+						for (const attachment of m.attachments) {
+							if (attachment.type === "text" && attachment.content) {
+								msgContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``;
+							}
+						}
+					}
+					return { role: m.role, content: msgContent };
+				}),
+			];
+
+			// Determine model
+			let modelToUse = this.selectedChatModel;
+			if (!modelToUse) {
+				for (const [, instanceWrapper] of Object.entries(this.instances)) {
+					if (instanceWrapper && typeof instanceWrapper === "object") {
+						const keys = Object.keys(
+							instanceWrapper as Record<string, unknown>,
+						);
+						if (keys.length === 1) {
+							const instance = (instanceWrapper as Record<string, unknown>)[
+								keys[0]
+							] as { shardAssignments?: { modelId?: string } };
+							if (instance?.shardAssignments?.modelId) {
+								modelToUse = instance.shardAssignments.modelId;
+								break;
+							}
+						}
+					}
+				}
+			}
+
+			if (!modelToUse) {
+				throw new Error("No model available");
+			}
+
+			// Start timing
+			const requestStartTime = performance.now();
+			let firstTokenTime: number | null = null;
+			let tokenCount = tokensToKeep.length;
+
+			const response = await fetch("/v1/chat/completions", {
+				method: "POST",
+				headers: { "Content-Type": "application/json" },
+				body: JSON.stringify({
+					model: modelToUse,
+					messages: apiMessages,
+					stream: true,
+					logprobs: true,
+					top_logprobs: 5,
+					continue_from_prefix: true,
+				}),
+				signal: controller.signal,
+			});
+
+			if (!response.ok) {
+				const errorText = await response.text();
+				throw new Error(`API error: ${response.status} - ${errorText}`);
+			}
+
+			const reader = response.body?.getReader();
+			if (!reader) throw new Error("No response body");
+
+			const decoder = new TextDecoder();
+			let fullContent = prefixText;
+			let buffer = "";
+			const collectedTokens: TokenData[] = [...tokensToKeep];
+
+			while (true) {
+				const { done, value } = await reader.read();
+				if (done) break;
+
+				buffer += decoder.decode(value, { stream: true });
+				const lines = buffer.split("\n");
+				buffer = lines.pop() || "";
+
+				for (const line of lines) {
+					const trimmed = line.trim();
+					if (!trimmed || trimmed === "data: [DONE]") continue;
+
+					if (trimmed.startsWith("data: ")) {
+						try {
+							const json = JSON.parse(trimmed.slice(6));
+							const delta = json.choices?.[0]?.delta?.content;
+							if (delta) {
+								if (firstTokenTime === null) {
+									firstTokenTime = performance.now();
+									this.ttftMs = firstTokenTime - requestStartTime;
+								}
+
+								tokenCount += 1;
+								this.totalTokens = tokenCount;
+
+								if (
+									firstTokenTime !== null &&
+									tokenCount > tokensToKeep.length
+								) {
+									const elapsed = performance.now() - firstTokenTime;
+									this.tps =
+										((tokenCount - tokensToKeep.length) / elapsed) * 1000;
+								}
+
+								// Extract logprobs
+								const logprobsData = json.choices?.[0]?.logprobs;
+								if (logprobsData?.content?.[0]) {
+									const logprobItem = logprobsData.content[0];
+									collectedTokens.push({
+										token: logprobItem.token || delta,
+										logprob: logprobItem.logprob ?? 0,
+										probability: Math.exp(logprobItem.logprob ?? 0),
+										topLogprobs: (logprobItem.top_logprobs || []).map(
+											(item: {
+												token: string;
+												logprob: number;
+												bytes?: number[];
+											}) => ({
+												token: item.token,
+												logprob: item.logprob,
+												bytes: item.bytes,
+											}),
+										),
+									});
+								}
+
+								fullContent += delta;
+								const { displayContent, thinkingContent } =
+									this.stripThinkingTags(fullContent);
+								this.currentResponse = displayContent;
+
+								this.messages[messageIdx].content = displayContent;
+								this.messages[messageIdx].thinking =
+									thinkingContent || undefined;
+								this.messages[messageIdx].tokens = [...collectedTokens];
+								this.persistActiveConversation();
+							}
+						} catch {
+							// Skip malformed JSON
+						}
+					}
+				}
+			}
+
+			// Final update
+			const { displayContent, thinkingContent } =
+				this.stripThinkingTags(fullContent);
+			this.messages[messageIdx].content = displayContent;
+			this.messages[messageIdx].thinking = thinkingContent || undefined;
+			this.messages[messageIdx].tokens = collectedTokens;
+
+			if (this.ttftMs !== null) {
+				this.messages[messageIdx].ttftMs = this.ttftMs;
+			}
+			if (this.tps !== null) {
+				this.messages[messageIdx].tps = this.tps;
+			}
+			this.persistActiveConversation();
+		} catch (error) {
+			if (error instanceof Error && error.name === "AbortError") {
+				return;
+			}
+			console.error("Error regenerating from token:", error);
+			this.messages[messageIdx].content =
+				`${prefixText}\n\nError: ${error instanceof Error ? error.message : "Unknown error"}`;
+			this.persistActiveConversation();
+		} finally {
+			if (this.currentRequestController === controller) {
+				this.currentRequestController = null;
+			}
 			this.isLoading = false;
 			this.currentResponse = "";
 			this.updateActiveConversation();
@@ -1650,6 +1965,8 @@ export const editMessage = (messageId: string, newContent: string) =>
 export const editAndRegenerate = (messageId: string, newContent: string) =>
 	appStore.editAndRegenerate(messageId, newContent);
 export const regenerateLastResponse = () => appStore.regenerateLastResponse();
+export const regenerateFromToken = (messageId: string, tokenIndex: number) =>
+	appStore.regenerateFromToken(messageId, tokenIndex);
 
 // Conversation actions
 export const conversations = () => appStore.conversations;

dashboard/src/routes/+page.svelte

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

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

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

pyproject.toml

@@ -17,8 +17,8 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx==0.30.3; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.3; sys_platform == 'linux'",
+    "mlx==0.30.1; sys_platform == 'darwin'",
+    "mlx[cpu]==0.30.1; sys_platform == 'linux'",
     "mlx-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",

src/exo/master/adapters/__init__.py

@@ -0,0 +1 @@
+"""API adapters for different API formats (Claude, OpenAI Responses, etc.)."""

src/exo/master/adapters/chat_completions.py

@@ -0,0 +1,186 @@
+"""OpenAI Chat Completions API adapter for converting requests/responses."""
+
+import time
+from collections.abc import AsyncGenerator
+
+from loguru import logger
+
+from exo.shared.types.api import (
+    ChatCompletionChoice,
+    ChatCompletionMessage,
+    ChatCompletionMessageText,
+    ChatCompletionResponse,
+    ChatCompletionTaskParams,
+    ErrorInfo,
+    ErrorResponse,
+    FinishReason,
+    Logprobs,
+    LogprobsContentItem,
+    StreamingChoiceResponse,
+)
+from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import ResponseInputMessage, ResponsesRequest
+
+
+def chat_request_to_internal(request: ChatCompletionTaskParams) -> ResponsesRequest:
+    """Convert Chat Completions API request to ResponsesRequest (canonical internal format).
+
+    Extracts system message as instructions, converts messages to input.
+    """
+    instructions: str | None = None
+    input_messages: list[ResponseInputMessage] = []
+
+    for msg in request.messages:
+        # Normalize content to string
+        content: str
+        if msg.content is None:
+            content = ""
+        elif isinstance(msg.content, str):
+            content = msg.content
+        elif isinstance(msg.content, ChatCompletionMessageText):
+            content = msg.content.text
+        else:
+            # List of ChatCompletionMessageText
+            content = "\n".join(item.text for item in msg.content)
+
+        # Extract system message as instructions
+        if msg.role == "system":
+            if instructions is None:
+                instructions = content
+            else:
+                # Append additional system messages
+                instructions = f"{instructions}\n{content}"
+        else:
+            # Convert to ResponseInputMessage (only user, assistant, developer roles)
+            if msg.role in ("user", "assistant", "developer"):
+                input_messages.append(
+                    ResponseInputMessage(role=msg.role, content=content)
+                )
+
+    return ResponsesRequest(
+        model=request.model,
+        input=input_messages if input_messages else "",
+        instructions=instructions,
+        max_output_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        top_k=request.top_k,
+        stop=request.stop,
+        seed=request.seed,
+        stream=request.stream,
+        tools=request.tools,
+        continue_from_prefix=request.continue_from_prefix,
+    )
+
+
+def chunk_to_response(
+    chunk: TokenChunk, command_id: CommandId
+) -> ChatCompletionResponse:
+    """Convert a TokenChunk to a streaming ChatCompletionResponse."""
+    # Build logprobs if available
+    logprobs: Logprobs | None = None
+    if chunk.logprob is not None:
+        logprobs = Logprobs(
+            content=[
+                LogprobsContentItem(
+                    token=chunk.text,
+                    logprob=chunk.logprob,
+                    top_logprobs=chunk.top_logprobs or [],
+                )
+            ]
+        )
+
+    return ChatCompletionResponse(
+        id=command_id,
+        created=int(time.time()),
+        model=chunk.model,
+        choices=[
+            StreamingChoiceResponse(
+                index=0,
+                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+                logprobs=logprobs,
+                finish_reason=chunk.finish_reason,
+            )
+        ],
+    )
+
+
+async def generate_chat_stream(
+    command_id: CommandId,
+    chunk_stream: AsyncGenerator[TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+    """Generate Chat Completions API streaming events from TokenChunks."""
+    try:
+        async for chunk in chunk_stream:
+            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"
+                logger.info(f"generate_chat_stream ending (error): {command_id}")
+                return
+
+            chunk_response = chunk_to_response(chunk, command_id)
+            yield f"data: {chunk_response.model_dump_json()}\n\n"
+
+            if chunk.finish_reason is not None:
+                logger.info(
+                    f"generate_chat_stream yielding [DONE] for finish_reason={chunk.finish_reason}: {command_id}"
+                )
+                yield "data: [DONE]\n\n"
+                logger.info(f"generate_chat_stream returning: {command_id}")
+                return
+    finally:
+        logger.info(f"generate_chat_stream finally block: {command_id}")
+
+
+async def collect_chat_response(
+    command_id: CommandId,
+    chunk_stream: AsyncGenerator[TokenChunk, None],
+) -> ChatCompletionResponse:
+    """Collect all token chunks and return a single ChatCompletionResponse."""
+    text_parts: list[str] = []
+    model: str | None = None
+    finish_reason: FinishReason | None = None
+    error_message: str | None = None
+
+    async for chunk in chunk_stream:
+        if chunk.finish_reason == "error":
+            error_message = chunk.error_message or "Internal server error"
+            break
+
+        if model is None:
+            model = chunk.model
+
+        text_parts.append(chunk.text)
+
+        if chunk.finish_reason is not None:
+            finish_reason = chunk.finish_reason
+
+    if error_message is not None:
+        raise ValueError(error_message)
+
+    combined_text = "".join(text_parts)
+    assert model is not None
+
+    return ChatCompletionResponse(
+        id=command_id,
+        created=int(time.time()),
+        model=model,
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                message=ChatCompletionMessage(
+                    role="assistant",
+                    content=combined_text,
+                ),
+                finish_reason=finish_reason,
+            )
+        ],
+    )

src/exo/master/adapters/claude.py

@@ -0,0 +1,190 @@
+"""Claude Messages API adapter for converting requests/responses."""
+
+from collections.abc import AsyncGenerator
+
+from exo.shared.types.api import FinishReason
+from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.claude_api import (
+    ClaudeContentBlockDeltaEvent,
+    ClaudeContentBlockStartEvent,
+    ClaudeContentBlockStopEvent,
+    ClaudeMessageDelta,
+    ClaudeMessageDeltaEvent,
+    ClaudeMessageDeltaUsage,
+    ClaudeMessagesRequest,
+    ClaudeMessagesResponse,
+    ClaudeMessageStart,
+    ClaudeMessageStartEvent,
+    ClaudeMessageStopEvent,
+    ClaudeStopReason,
+    ClaudeTextBlock,
+    ClaudeTextDelta,
+    ClaudeUsage,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import ResponseInputMessage, ResponsesRequest
+
+
+def finish_reason_to_claude_stop_reason(
+    finish_reason: FinishReason | None,
+) -> ClaudeStopReason | None:
+    """Map OpenAI finish_reason to Claude stop_reason."""
+    if finish_reason is None:
+        return None
+    mapping: dict[FinishReason, ClaudeStopReason] = {
+        "stop": "end_turn",
+        "length": "max_tokens",
+        "tool_calls": "tool_use",
+        "content_filter": "end_turn",
+        "function_call": "tool_use",
+    }
+    return mapping.get(finish_reason, "end_turn")
+
+
+def claude_request_to_internal(request: ClaudeMessagesRequest) -> ResponsesRequest:
+    """Convert Claude Messages API request to ResponsesRequest (canonical internal format).
+
+    Converts Claude's system parameter to instructions,
+    and messages to input.
+    """
+    # Handle system message
+    instructions: str | None = None
+    if request.system:
+        if isinstance(request.system, str):
+            instructions = request.system
+        else:
+            # List of text blocks
+            instructions = "".join(block.text for block in request.system)
+
+    # Convert messages to input
+    input_messages: list[ResponseInputMessage] = []
+    for msg in request.messages:
+        content: str
+        if isinstance(msg.content, str):
+            content = msg.content
+        else:
+            # Concatenate text blocks (images not supported for MVP)
+            text_parts: list[str] = []
+            for block in msg.content:
+                if isinstance(block, ClaudeTextBlock):
+                    text_parts.append(block.text)
+            content = "".join(text_parts)
+
+        # Claude uses "user" and "assistant" roles
+        input_messages.append(ResponseInputMessage(role=msg.role, content=content))
+
+    return ResponsesRequest(
+        model=request.model,
+        input=input_messages if input_messages else "",
+        instructions=instructions,
+        max_output_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        top_k=request.top_k,
+        stop=request.stop_sequences,
+        stream=request.stream,
+    )
+
+
+async def collect_claude_response(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[TokenChunk, None],
+) -> ClaudeMessagesResponse:
+    """Collect all token chunks and return a single ClaudeMessagesResponse."""
+    text_parts: list[str] = []
+    stop_reason: ClaudeStopReason | None = None
+    last_stats = None
+    error_message: str | None = None
+
+    async for chunk in chunk_stream:
+        if chunk.finish_reason == "error":
+            error_message = chunk.error_message or "Internal server error"
+            break
+
+        text_parts.append(chunk.text)
+        last_stats = chunk.stats or last_stats
+
+        if chunk.finish_reason is not None:
+            stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
+
+    if error_message is not None:
+        raise ValueError(error_message)
+
+    combined_text = "".join(text_parts)
+
+    # Use actual usage data from stats if available
+    input_tokens = last_stats.prompt_tokens if last_stats else 0
+    output_tokens = last_stats.generation_tokens if last_stats else 0
+
+    return ClaudeMessagesResponse(
+        id=f"msg_{command_id}",
+        model=model,
+        content=[ClaudeTextBlock(text=combined_text)],
+        stop_reason=stop_reason,
+        usage=ClaudeUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+        ),
+    )
+
+
+async def generate_claude_stream(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+    """Generate Claude Messages API streaming events from TokenChunks."""
+    # Initial message_start event
+    initial_message = ClaudeMessageStart(
+        id=f"msg_{command_id}",
+        model=model,
+        content=[],
+        stop_reason=None,
+        usage=ClaudeUsage(input_tokens=0, output_tokens=0),
+    )
+    start_event = ClaudeMessageStartEvent(message=initial_message)
+    yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
+
+    # content_block_start
+    block_start = ClaudeContentBlockStartEvent(
+        index=0, content_block=ClaudeTextBlock(text="")
+    )
+    yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+    output_tokens = 0
+    stop_reason: ClaudeStopReason | None = None
+    last_stats = None
+
+    async for chunk in chunk_stream:
+        output_tokens += 1  # Count each chunk as one token
+        last_stats = chunk.stats or last_stats
+
+        # content_block_delta
+        delta_event = ClaudeContentBlockDeltaEvent(
+            index=0,
+            delta=ClaudeTextDelta(text=chunk.text),
+        )
+        yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+
+        if chunk.finish_reason is not None:
+            stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
+
+    # Use actual token count from stats if available
+    if last_stats is not None:
+        output_tokens = last_stats.generation_tokens
+
+    # content_block_stop
+    block_stop = ClaudeContentBlockStopEvent(index=0)
+    yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+    # message_delta
+    message_delta = ClaudeMessageDeltaEvent(
+        delta=ClaudeMessageDelta(stop_reason=stop_reason),
+        usage=ClaudeMessageDeltaUsage(output_tokens=output_tokens),
+    )
+    yield f"event: message_delta\ndata: {message_delta.model_dump_json()}\n\n"
+
+    # message_stop
+    message_stop = ClaudeMessageStopEvent()
+    yield f"event: message_stop\ndata: {message_stop.model_dump_json()}\n\n"

src/exo/master/adapters/responses.py

@@ -0,0 +1,173 @@
+"""OpenAI Responses API adapter for converting requests/responses.
+
+ResponsesRequest is the canonical internal format. Responses API is the most featureful,
+making it the natural choice for the internal format. All other API formats (Chat
+Completions, Claude) are converted TO ResponsesRequest.
+"""
+
+from collections.abc import AsyncGenerator
+
+from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import (
+    ResponseCompletedEvent,
+    ResponseContentPartAddedEvent,
+    ResponseContentPartDoneEvent,
+    ResponseCreatedEvent,
+    ResponseInProgressEvent,
+    ResponseMessageItem,
+    ResponseOutputItemAddedEvent,
+    ResponseOutputItemDoneEvent,
+    ResponseOutputText,
+    ResponsesResponse,
+    ResponseTextDeltaEvent,
+    ResponseTextDoneEvent,
+    ResponseUsage,
+)
+
+
+async def collect_responses_response(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[TokenChunk, None],
+) -> ResponsesResponse:
+    """Collect all token chunks and return a single ResponsesResponse."""
+    response_id = f"resp_{command_id}"
+    item_id = f"item_{command_id}"
+    accumulated_text = ""
+    last_stats = None
+    error_message: str | None = None
+
+    async for chunk in chunk_stream:
+        if chunk.finish_reason == "error":
+            error_message = chunk.error_message or "Internal server error"
+            break
+
+        accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
+
+    if error_message is not None:
+        raise ValueError(error_message)
+
+    # Create usage from stats if available
+    usage = None
+    if last_stats is not None:
+        usage = ResponseUsage(
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
+        )
+
+    output_item = ResponseMessageItem(
+        id=item_id,
+        content=[ResponseOutputText(text=accumulated_text)],
+        status="completed",
+    )
+
+    return ResponsesResponse(
+        id=response_id,
+        model=model,
+        status="completed",
+        output=[output_item],
+        output_text=accumulated_text,
+        usage=usage,
+    )
+
+
+async def generate_responses_stream(
+    command_id: CommandId,
+    model: str,
+    chunk_stream: AsyncGenerator[TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+    """Generate OpenAI Responses API streaming events from TokenChunks."""
+    response_id = f"resp_{command_id}"
+    item_id = f"item_{command_id}"
+
+    # response.created
+    initial_response = ResponsesResponse(
+        id=response_id,
+        model=model,
+        status="in_progress",
+        output=[],
+        output_text="",
+    )
+    created_event = ResponseCreatedEvent(response=initial_response)
+    yield f"event: response.created\ndata: {created_event.model_dump_json()}\n\n"
+
+    # response.in_progress
+    in_progress_event = ResponseInProgressEvent(response=initial_response)
+    yield f"event: response.in_progress\ndata: {in_progress_event.model_dump_json()}\n\n"
+
+    # response.output_item.added
+    initial_item = ResponseMessageItem(
+        id=item_id,
+        content=[ResponseOutputText(text="")],
+        status="in_progress",
+    )
+    item_added = ResponseOutputItemAddedEvent(output_index=0, item=initial_item)
+    yield f"event: response.output_item.added\ndata: {item_added.model_dump_json()}\n\n"
+
+    # response.content_part.added
+    initial_part = ResponseOutputText(text="")
+    part_added = ResponseContentPartAddedEvent(
+        output_index=0, content_index=0, part=initial_part
+    )
+    yield f"event: response.content_part.added\ndata: {part_added.model_dump_json()}\n\n"
+
+    accumulated_text = ""
+    last_stats = None
+
+    async for chunk in chunk_stream:
+        accumulated_text += chunk.text
+        last_stats = chunk.stats or last_stats
+
+        # response.output_text.delta
+        delta_event = ResponseTextDeltaEvent(
+            output_index=0,
+            content_index=0,
+            delta=chunk.text,
+        )
+        yield f"event: response.output_text.delta\ndata: {delta_event.model_dump_json()}\n\n"
+
+    # response.output_text.done
+    text_done = ResponseTextDoneEvent(
+        output_index=0, content_index=0, text=accumulated_text
+    )
+    yield f"event: response.output_text.done\ndata: {text_done.model_dump_json()}\n\n"
+
+    # response.content_part.done
+    final_part = ResponseOutputText(text=accumulated_text)
+    part_done = ResponseContentPartDoneEvent(
+        output_index=0, content_index=0, part=final_part
+    )
+    yield f"event: response.content_part.done\ndata: {part_done.model_dump_json()}\n\n"
+
+    # response.output_item.done
+    final_item = ResponseMessageItem(
+        id=item_id,
+        content=[ResponseOutputText(text=accumulated_text)],
+        status="completed",
+    )
+    item_done = ResponseOutputItemDoneEvent(output_index=0, item=final_item)
+    yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
+
+    # Create usage from stats if available
+    usage = None
+    if last_stats is not None:
+        usage = ResponseUsage(
+            input_tokens=last_stats.prompt_tokens,
+            output_tokens=last_stats.generation_tokens,
+            total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
+        )
+
+    # response.completed
+    final_response = ResponsesResponse(
+        id=response_id,
+        model=model,
+        status="completed",
+        output=[final_item],
+        output_text=accumulated_text,
+        usage=usage,
+    )
+    completed_event = ResponseCompletedEvent(response=final_response)
+    yield f"event: response.completed\ndata: {completed_event.model_dump_json()}\n\n"

src/exo/master/api.py

@@ -15,18 +15,33 @@ from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger
 
+from exo.master.adapters.chat_completions import (
+    chat_request_to_internal,
+    collect_chat_response,
+    generate_chat_stream,
+)
+from exo.master.adapters.claude import (
+    claude_request_to_internal,
+    collect_claude_response,
+    generate_claude_stream,
+)
+from exo.master.adapters.responses import (
+    collect_responses_response,
+    generate_responses_stream,
+)
 from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
-from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
-from exo.shared.models.model_meta import get_model_card
+from exo.shared.models.model_cards import MODEL_CARDS
+from exo.shared.models.model_meta import get_model_meta
 from exo.shared.types.api import (
     BenchChatCompletionResponse,
     BenchChatCompletionTaskParams,
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionResponse,
+    ChatCompletionTaskParams,
     CreateInstanceParams,
     CreateInstanceResponse,
     DeleteInstanceResponse,
@@ -39,9 +54,12 @@ from exo.shared.types.api import (
     PlaceInstanceParams,
     PlacementPreview,
     PlacementPreviewResponse,
-    StreamingChoiceResponse,
 )
 from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.claude_api import (
+    ClaudeMessagesRequest,
+    ClaudeMessagesResponse,
+)
 from exo.shared.types.commands import (
     ChatCompletion,
     Command,
@@ -59,8 +77,12 @@ from exo.shared.types.events import (
     IndexedEvent,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.openai_responses import (
+    ResponsesRequest,
+    ResponsesResponse,
+)
 from exo.shared.types.state import State
-from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
@@ -69,29 +91,12 @@ from exo.utils.dashboard_path import find_dashboard
 from exo.utils.event_buffer import OrderedBuffer
 
 
-def chunk_to_response(
-    chunk: TokenChunk, command_id: CommandId
-) -> ChatCompletionResponse:
-    return ChatCompletionResponse(
-        id=command_id,
-        created=int(time.time()),
-        model=chunk.model,
-        choices=[
-            StreamingChoiceResponse(
-                index=0,
-                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
-                finish_reason=chunk.finish_reason,
-            )
-        ],
-    )
-
-
-async def resolve_model_card(model_id: str) -> ModelCard:
+async def resolve_model_meta(model_id: str) -> ModelMetadata:
     if model_id in MODEL_CARDS:
         model_card = MODEL_CARDS[model_id]
-        return model_card
+        return model_card.metadata
     else:
-        return await get_model_card(model_id)
+        return await get_model_meta(model_id)
 
 
 class API:
@@ -191,12 +196,14 @@ class API:
             self.chat_completions
         )
         self.app.post("/bench/chat/completions")(self.bench_chat_completions)
+        self.app.post("/v1/messages", response_model=None)(self.claude_messages)
+        self.app.post("/v1/responses", response_model=None)(self.openai_responses)
         self.app.get("/state")(lambda: self.state)
         self.app.get("/events")(lambda: self._event_log)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
-            model_card=await resolve_model_card(payload.model_id),
+            model_meta=await resolve_model_meta(payload.model_id),
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
@@ -206,15 +213,15 @@ class API:
         return CreateInstanceResponse(
             message="Command received.",
             command_id=command.command_id,
-            model_card=command.model_card,
+            model_meta=command.model_meta,
         )
 
     async def create_instance(
         self, payload: CreateInstanceParams
     ) -> CreateInstanceResponse:
         instance = payload.instance
-        model_card = await resolve_model_card(instance.shard_assignments.model_id)
-        required_memory = model_card.storage_size
+        model_meta = await resolve_model_meta(instance.shard_assignments.model_id)
+        required_memory = model_meta.storage_size
         available_memory = self._calculate_total_available_memory()
 
         if required_memory > available_memory:
@@ -231,7 +238,7 @@ class API:
         return CreateInstanceResponse(
             message="Command received.",
             command_id=command.command_id,
-            model_card=model_card,
+            model_meta=model_meta,
         )
 
     async def get_placement(
@@ -241,18 +248,16 @@ class API:
         instance_meta: InstanceMeta = InstanceMeta.MlxRing,
         min_nodes: int = 1,
     ) -> Instance:
-        model_card = await resolve_model_card(model_id)
+        model_meta = await resolve_model_meta(model_id)
 
         try:
             placements = get_instance_placements(
                 PlaceInstance(
-                    model_card=model_card,
+                    model_meta=model_meta,
                     sharding=sharding,
                     instance_meta=instance_meta,
                     min_nodes=min_nodes,
                 ),
-                node_memory=self.state.node_memory,
-                node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
             )
@@ -279,7 +284,7 @@ class API:
         if len(list(self.state.topology.list_nodes())) == 0:
             return PlacementPreviewResponse(previews=[])
 
-        cards = [card for card in MODEL_CARDS.values() if card.model_id == model_id]
+        cards = [card for card in MODEL_CARDS.values() if card.short_id == model_id]
         if not cards:
             raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
 
@@ -297,33 +302,32 @@ class API:
         # TODO: PDD
         # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
 
-        for model_card in cards:
+        for card in cards:
+            model_meta = card.metadata
             for sharding, instance_meta, min_nodes in instance_combinations:
                 try:
                     placements = get_instance_placements(
                         PlaceInstance(
-                            model_card=model_card,
+                            model_meta=model_meta,
                             sharding=sharding,
                             instance_meta=instance_meta,
                             min_nodes=min_nodes,
                         ),
-                        node_memory=self.state.node_memory,
-                        node_network=self.state.node_network,
                         topology=self.state.topology,
                         current_instances=self.state.instances,
                     )
                 except ValueError as exc:
-                    if (model_card.model_id, sharding, instance_meta, 0) not in seen:
+                    if (card.model_id, sharding, instance_meta, 0) not in seen:
                         previews.append(
                             PlacementPreview(
-                                model_id=model_card.model_id,
+                                model_id=card.model_id,
                                 sharding=sharding,
                                 instance_meta=instance_meta,
                                 instance=None,
                                 error=str(exc),
                             )
                         )
-                    seen.add((model_card.model_id, sharding, instance_meta, 0))
+                    seen.add((card.model_id, sharding, instance_meta, 0))
                     continue
 
                 current_ids = set(self.state.instances.keys())
@@ -334,17 +338,17 @@ class API:
                 ]
 
                 if len(new_instances) != 1:
-                    if (model_card.model_id, sharding, instance_meta, 0) not in seen:
+                    if (card.model_id, sharding, instance_meta, 0) not in seen:
                         previews.append(
                             PlacementPreview(
-                                model_id=model_card.model_id,
+                                model_id=card.model_id,
                                 sharding=sharding,
                                 instance_meta=instance_meta,
                                 instance=None,
                                 error="Expected exactly one new instance from placement",
                             )
                         )
-                    seen.add((model_card.model_id, sharding, instance_meta, 0))
+                    seen.add((card.model_id, sharding, instance_meta, 0))
                     continue
 
                 instance = new_instances[0]
@@ -353,7 +357,7 @@ class API:
 
                 memory_delta_by_node: dict[str, int] = {}
                 if node_ids:
-                    total_bytes = model_card.storage_size.in_bytes
+                    total_bytes = model_meta.storage_size.in_bytes
                     per_node = total_bytes // len(node_ids)
                     remainder = total_bytes % len(node_ids)
                     for index, node_id in enumerate(sorted(node_ids, key=str)):
@@ -361,14 +365,14 @@ class API:
                         memory_delta_by_node[str(node_id)] = per_node + extra
 
                 if (
-                    model_card.model_id,
+                    card.model_id,
                     sharding,
                     instance_meta,
                     len(node_ids),
                 ) not in seen:
                     previews.append(
                         PlacementPreview(
-                            model_id=model_card.model_id,
+                            model_id=card.model_id,
                             sharding=sharding,
                             instance_meta=instance_meta,
                             instance=instance,
@@ -376,7 +380,7 @@ class API:
                             error=None,
                         )
                     )
-                seen.add((model_card.model_id, sharding, instance_meta, len(node_ids)))
+                seen.add((card.model_id, sharding, instance_meta, len(node_ids)))
 
         return PlacementPreviewResponse(previews=previews)
 
@@ -399,11 +403,13 @@ class API:
             instance_id=instance_id,
         )
 
-    async def _chat_chunk_stream(
+    async def _token_chunk_stream(
         self, command_id: CommandId
     ) -> AsyncGenerator[TokenChunk, None]:
-        """Yield `TokenChunk`s for a given command until completion."""
+        """Yield `TokenChunk`s for a given command until completion.
 
+        This is the internal low-level stream used by all API adapters.
+        """
         try:
             self._chat_completion_queues[command_id], recv = channel[TokenChunk]()
 
@@ -426,77 +432,6 @@ class API:
             await self._send(command)
             del self._chat_completion_queues[command_id]
 
-    async def _generate_chat_stream(
-        self, command_id: CommandId
-    ) -> AsyncGenerator[str, None]:
-        """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
-            )
-            logger.debug(f"chunk_response: {chunk_response}")
-
-            yield f"data: {chunk_response.model_dump_json()}\n\n"
-
-            if chunk.finish_reason is not None:
-                yield "data: [DONE]\n\n"
-
-    async def _collect_chat_completion(
-        self, command_id: CommandId
-    ) -> ChatCompletionResponse:
-        """Collect all token chunks for a chat completion and return a single response."""
-
-        text_parts: list[str] = []
-        model: str | None = None
-        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
-
-            text_parts.append(chunk.text)
-
-            if chunk.finish_reason is not None:
-                finish_reason = chunk.finish_reason
-
-        combined_text = "".join(text_parts)
-        assert model is not None
-
-        return ChatCompletionResponse(
-            id=command_id,
-            created=int(time.time()),
-            model=model,
-            choices=[
-                ChatCompletionChoice(
-                    index=0,
-                    message=ChatCompletionMessage(
-                        role="assistant",
-                        content=combined_text,
-                    ),
-                    finish_reason=finish_reason,
-                )
-            ],
-        )
-
     async def _collect_chat_completion_with_stats(
         self, command_id: CommandId
     ) -> BenchChatCompletionResponse:
@@ -506,7 +441,7 @@ class API:
 
         stats: GenerationStats | None = None
 
-        async for chunk in self._chat_chunk_stream(command_id):
+        async for chunk in self._token_chunk_stream(command_id):
             if chunk.finish_reason == "error":
                 raise HTTPException(
                     status_code=500,
@@ -550,60 +485,169 @@ class API:
     async def chat_completions(
         self, payload: ChatCompletionTaskParams
     ) -> ChatCompletionResponse | StreamingResponse:
-        """Handle chat completions, supporting both streaming and non-streaming responses."""
-        model_card = await resolve_model_card(payload.model)
-        payload.model = model_card.model_id
+        """OpenAI Chat Completions API - adapter."""
+        internal_params = chat_request_to_internal(payload)
+        model_meta = await resolve_model_meta(internal_params.model)
+        internal_params.model = model_meta.model_id
 
         if not any(
-            instance.shard_assignments.model_id == payload.model
+            instance.shard_assignments.model_id == internal_params.model
             for instance in self.state.instances.values()
         ):
-            await self._trigger_notify_user_to_download_model(payload.model)
+            await self._trigger_notify_user_to_download_model(internal_params.model)
             raise HTTPException(
-                status_code=404, detail=f"No instance found for model {payload.model}"
+                status_code=404,
+                detail=f"No instance found for model {internal_params.model}",
             )
 
-        command = ChatCompletion(
-            request_params=payload,
-        )
+        command = ChatCompletion(request_params=internal_params)
         await self._send(command)
+
         if payload.stream:
             return StreamingResponse(
-                self._generate_chat_stream(command.command_id),
+                generate_chat_stream(
+                    command.command_id,
+                    self._token_chunk_stream(command.command_id),
+                ),
                 media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "close",
+                    "X-Accel-Buffering": "no",
+                },
             )
 
-        return await self._collect_chat_completion(command.command_id)
+        try:
+            return await collect_chat_response(
+                command.command_id,
+                self._token_chunk_stream(command.command_id),
+            )
+        except ValueError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
 
     async def bench_chat_completions(
         self, payload: BenchChatCompletionTaskParams
     ) -> BenchChatCompletionResponse:
-        model_card = await resolve_model_card(payload.model)
-        payload.model = model_card.model_id
+        # Convert to internal format (BenchChatCompletionTaskParams extends ChatCompletionTaskParams)
+        internal_params = chat_request_to_internal(payload)
+        model_meta = await resolve_model_meta(internal_params.model)
+        internal_params.model = model_meta.model_id
 
         if not any(
-            instance.shard_assignments.model_id == payload.model
+            instance.shard_assignments.model_id == internal_params.model
             for instance in self.state.instances.values()
         ):
-            await self._trigger_notify_user_to_download_model(payload.model)
+            await self._trigger_notify_user_to_download_model(internal_params.model)
             raise HTTPException(
-                status_code=404, detail=f"No instance found for model {payload.model}"
+                status_code=404,
+                detail=f"No instance found for model {internal_params.model}",
             )
 
-        payload.stream = False
+        internal_params.stream = False
 
-        command = ChatCompletion(request_params=payload)
+        command = ChatCompletion(request_params=internal_params)
         await self._send(command)
 
         response = await self._collect_chat_completion_with_stats(command.command_id)
         return response
 
+    async def claude_messages(
+        self, payload: ClaudeMessagesRequest
+    ) -> ClaudeMessagesResponse | StreamingResponse:
+        """Claude Messages API - adapter."""
+        internal_params = claude_request_to_internal(payload)
+        model_meta = await resolve_model_meta(internal_params.model)
+        internal_params.model = model_meta.model_id
+
+        if not any(
+            instance.shard_assignments.model_id == internal_params.model
+            for instance in self.state.instances.values()
+        ):
+            await self._trigger_notify_user_to_download_model(internal_params.model)
+            raise HTTPException(
+                status_code=404,
+                detail=f"No instance found for model {internal_params.model}",
+            )
+
+        command = ChatCompletion(request_params=internal_params)
+        await self._send(command)
+
+        if payload.stream:
+            return StreamingResponse(
+                generate_claude_stream(
+                    command.command_id,
+                    payload.model,
+                    self._token_chunk_stream(command.command_id),
+                ),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "close",
+                    "X-Accel-Buffering": "no",
+                },
+            )
+
+        try:
+            return await collect_claude_response(
+                command.command_id,
+                payload.model,
+                self._token_chunk_stream(command.command_id),
+            )
+        except ValueError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    async def openai_responses(
+        self, payload: ResponsesRequest
+    ) -> ResponsesResponse | StreamingResponse:
+        """OpenAI Responses API - native format (no conversion needed)."""
+        model_meta = await resolve_model_meta(payload.model)
+        # Update model to resolved model_id
+        request_params = payload.model_copy(update={"model": model_meta.model_id})
+
+        if not any(
+            instance.shard_assignments.model_id == request_params.model
+            for instance in self.state.instances.values()
+        ):
+            await self._trigger_notify_user_to_download_model(request_params.model)
+            raise HTTPException(
+                status_code=404,
+                detail=f"No instance found for model {request_params.model}",
+            )
+
+        command = ChatCompletion(request_params=request_params)
+        await self._send(command)
+
+        if payload.stream:
+            return StreamingResponse(
+                generate_responses_stream(
+                    command.command_id,
+                    payload.model,
+                    self._token_chunk_stream(command.command_id),
+                ),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "close",
+                    "X-Accel-Buffering": "no",
+                },
+            )
+
+        try:
+            return await collect_responses_response(
+                command.command_id,
+                payload.model,
+                self._token_chunk_stream(command.command_id),
+            )
+        except ValueError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
     def _calculate_total_available_memory(self) -> Memory:
         """Calculate total available memory across all nodes in bytes."""
         total_available = Memory()
 
-        for memory in self.state.node_memory.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
 
@@ -612,13 +656,13 @@ class API:
         return ModelList(
             data=[
                 ModelListModel(
-                    id=card.model_id,
+                    id=card.short_id,
                     hugging_face_id=card.model_id,
-                    name=card.model_id.short(),
-                    description="",
-                    tags=[],
-                    storage_size_megabytes=int(card.storage_size.in_mb),
-                    supports_tensor=card.supports_tensor,
+                    name=card.name,
+                    description=card.description,
+                    tags=card.tags,
+                    storage_size_megabytes=int(card.metadata.storage_size.in_mb),
+                    supports_tensor=card.metadata.supports_tensor,
                 )
                 for card in MODEL_CARDS.values()
             ]

src/exo/master/main.py

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

src/exo/master/placement.py

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

src/exo/master/placement_utils.py

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

src/exo/master/tests/conftest.py

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

src/exo/master/tests/test_claude_api.py

@@ -0,0 +1,283 @@
+"""Tests for Claude Messages API conversion functions and types."""
+
+import json
+from typing import Any, cast
+
+import pydantic
+import pytest
+
+from exo.master.adapters.claude import (
+    claude_request_to_internal,
+    finish_reason_to_claude_stop_reason,
+)
+from exo.shared.types.claude_api import (
+    ClaudeContentBlockDeltaEvent,
+    ClaudeContentBlockStartEvent,
+    ClaudeContentBlockStopEvent,
+    ClaudeMessage,
+    ClaudeMessageDelta,
+    ClaudeMessageDeltaEvent,
+    ClaudeMessageDeltaUsage,
+    ClaudeMessagesRequest,
+    ClaudeMessageStart,
+    ClaudeMessageStartEvent,
+    ClaudeMessageStopEvent,
+    ClaudeTextBlock,
+    ClaudeTextDelta,
+    ClaudeUsage,
+)
+
+
+class TestFinishReasonToClaudeStopReason:
+    """Tests for finish_reason to Claude stop_reason mapping."""
+
+    def test_stop_maps_to_end_turn(self):
+        assert finish_reason_to_claude_stop_reason("stop") == "end_turn"
+
+    def test_length_maps_to_max_tokens(self):
+        assert finish_reason_to_claude_stop_reason("length") == "max_tokens"
+
+    def test_tool_calls_maps_to_tool_use(self):
+        assert finish_reason_to_claude_stop_reason("tool_calls") == "tool_use"
+
+    def test_function_call_maps_to_tool_use(self):
+        assert finish_reason_to_claude_stop_reason("function_call") == "tool_use"
+
+    def test_content_filter_maps_to_end_turn(self):
+        assert finish_reason_to_claude_stop_reason("content_filter") == "end_turn"
+
+    def test_none_returns_none(self):
+        assert finish_reason_to_claude_stop_reason(None) is None
+
+
+class TestClaudeRequestToInternal:
+    """Tests for converting Claude Messages API requests to ResponsesRequest."""
+
+    def test_basic_request_conversion(self):
+        request = ClaudeMessagesRequest(
+            model="claude-3-opus",
+            max_tokens=100,
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+            ],
+        )
+        params = claude_request_to_internal(request)
+
+        assert params.model == "claude-3-opus"
+        assert params.max_output_tokens == 100
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+        assert params.input[0].role == "user"
+        assert params.input[0].content == "Hello"
+        assert params.instructions is None
+
+    def test_request_with_system_string(self):
+        request = ClaudeMessagesRequest(
+            model="claude-3-opus",
+            max_tokens=100,
+            system="You are a helpful assistant.",
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+            ],
+        )
+        params = claude_request_to_internal(request)
+
+        assert params.instructions == "You are a helpful assistant."
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+        assert params.input[0].role == "user"
+        assert params.input[0].content == "Hello"
+
+    def test_request_with_system_text_blocks(self):
+        request = ClaudeMessagesRequest(
+            model="claude-3-opus",
+            max_tokens=100,
+            system=[
+                ClaudeTextBlock(text="You are helpful. "),
+                ClaudeTextBlock(text="Be concise."),
+            ],
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+            ],
+        )
+        params = claude_request_to_internal(request)
+
+        assert params.instructions == "You are helpful. Be concise."
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+
+    def test_request_with_content_blocks(self):
+        request = ClaudeMessagesRequest(
+            model="claude-3-opus",
+            max_tokens=100,
+            messages=[
+                ClaudeMessage(
+                    role="user",
+                    content=[
+                        ClaudeTextBlock(text="First part. "),
+                        ClaudeTextBlock(text="Second part."),
+                    ],
+                ),
+            ],
+        )
+        params = claude_request_to_internal(request)
+
+        assert isinstance(params.input, list)
+        assert len(params.input) == 1
+        assert params.input[0].content == "First part. Second part."
+
+    def test_request_with_multi_turn_conversation(self):
+        request = ClaudeMessagesRequest(
+            model="claude-3-opus",
+            max_tokens=100,
+            messages=[
+                ClaudeMessage(role="user", content="Hello"),
+                ClaudeMessage(role="assistant", content="Hi there!"),
+                ClaudeMessage(role="user", content="How are you?"),
+            ],
+        )
+        params = claude_request_to_internal(request)
+
+        assert isinstance(params.input, list)
+        assert len(params.input) == 3
+        assert params.input[0].role == "user"
+        assert params.input[1].role == "assistant"
+        assert params.input[2].role == "user"
+
+    def test_request_with_optional_parameters(self):
+        request = ClaudeMessagesRequest(
+            model="claude-3-opus",
+            max_tokens=100,
+            messages=[ClaudeMessage(role="user", content="Hello")],
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            stop_sequences=["STOP", "END"],
+            stream=True,
+        )
+        params = claude_request_to_internal(request)
+
+        assert params.temperature == 0.7
+        assert params.top_p == 0.9
+        assert params.top_k == 40
+        assert params.stop == ["STOP", "END"]
+        assert params.stream is True
+
+
+class TestClaudeMessagesRequestValidation:
+    """Tests for Claude Messages API request validation."""
+
+    def test_request_requires_model(self):
+        with pytest.raises(pydantic.ValidationError):
+            ClaudeMessagesRequest.model_validate(
+                {
+                    "max_tokens": 100,
+                    "messages": [{"role": "user", "content": "Hello"}],
+                }
+            )
+
+    def test_request_requires_max_tokens(self):
+        with pytest.raises(pydantic.ValidationError):
+            ClaudeMessagesRequest.model_validate(
+                {
+                    "model": "claude-3-opus",
+                    "messages": [{"role": "user", "content": "Hello"}],
+                }
+            )
+
+    def test_request_requires_messages(self):
+        with pytest.raises(pydantic.ValidationError):
+            ClaudeMessagesRequest.model_validate(
+                {
+                    "model": "claude-3-opus",
+                    "max_tokens": 100,
+                }
+            )
+
+
+class TestClaudeStreamingEvents:
+    """Tests for Claude Messages API streaming event serialization."""
+
+    def test_message_start_event_format(self):
+        message = ClaudeMessageStart(
+            id="msg_123",
+            model="claude-3-opus",
+            content=[],
+            stop_reason=None,
+            usage=ClaudeUsage(input_tokens=10, output_tokens=0),
+        )
+        event = ClaudeMessageStartEvent(message=message)
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "message_start"
+        assert parsed["message"]["id"] == "msg_123"
+        assert parsed["message"]["type"] == "message"
+        assert parsed["message"]["role"] == "assistant"
+        assert parsed["message"]["model"] == "claude-3-opus"
+
+    def test_content_block_start_event_format(self):
+        event = ClaudeContentBlockStartEvent(
+            index=0,
+            content_block=ClaudeTextBlock(text=""),
+        )
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "content_block_start"
+        assert parsed["index"] == 0
+        assert parsed["content_block"]["type"] == "text"
+        assert parsed["content_block"]["text"] == ""
+
+    def test_content_block_delta_event_format(self):
+        event = ClaudeContentBlockDeltaEvent(
+            index=0,
+            delta=ClaudeTextDelta(text="Hello"),
+        )
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "content_block_delta"
+        assert parsed["index"] == 0
+        assert parsed["delta"]["type"] == "text_delta"
+        assert parsed["delta"]["text"] == "Hello"
+
+    def test_content_block_stop_event_format(self):
+        event = ClaudeContentBlockStopEvent(index=0)
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "content_block_stop"
+        assert parsed["index"] == 0
+
+    def test_message_delta_event_format(self):
+        event = ClaudeMessageDeltaEvent(
+            delta=ClaudeMessageDelta(stop_reason="end_turn"),
+            usage=ClaudeMessageDeltaUsage(output_tokens=25),
+        )
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "message_delta"
+        assert parsed["delta"]["stop_reason"] == "end_turn"
+        assert parsed["usage"]["output_tokens"] == 25
+
+    def test_message_stop_event_format(self):
+        event = ClaudeMessageStopEvent()
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "message_stop"
+
+    def test_sse_format(self):
+        """Test that SSE format is correctly generated."""
+        event = ClaudeContentBlockDeltaEvent(
+            index=0,
+            delta=ClaudeTextDelta(text="Hello"),
+        )
+        # Simulate the SSE format used in the streaming generator
+        sse_line = f"event: content_block_delta\ndata: {event.model_dump_json()}\n\n"
+
+        assert sse_line.startswith("event: content_block_delta\n")
+        assert "data: " in sse_line
+        assert sse_line.endswith("\n\n")

src/exo/master/tests/test_master.py

@@ -7,8 +7,6 @@ from loguru import logger
 
 from exo.master.main import Master
 from exo.routing.router import get_node_id_keypair
-from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
 from exo.shared.types.commands import (
     ChatCompletion,
     CommandId,
@@ -20,12 +18,16 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InstanceCreated,
-    NodeGatheredInfo,
+    NodePerformanceMeasured,
     TaskCreated,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.profiling import (
-    MemoryUsage,
+    MemoryPerformanceProfile,
+    NodePerformanceProfile,
+    SystemPerformanceProfile,
 )
 from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
 from exo.shared.types.tasks import TaskStatus
@@ -73,22 +75,29 @@ async def test_master():
         tg.start_soon(master.run)
 
         sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
-        # inject a NodeGatheredInfo event
-        logger.info("inject a NodeGatheredInfo 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=(
-                    NodeGatheredInfo(
+                    NodePerformanceMeasured(
                         when=str(datetime.now(tz=timezone.utc)),
                         node_id=node_id,
-                        info=MemoryUsage(
-                            ram_total=Memory.from_bytes(678948 * 1024),
-                            ram_available=Memory.from_bytes(678948 * 1024),
-                            swap_total=Memory.from_bytes(0),
-                            swap_available=Memory.from_bytes(0),
+                        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(),
                         ),
                     )
                 ),
@@ -99,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_memory) == 0:
+        while len(master.state.node_profiles) == 0:
             await anyio.sleep(0.001)
 
         logger.info("inject a CreateInstance Command")
@@ -109,8 +118,9 @@ async def test_master():
                 command=(
                     PlaceInstance(
                         command_id=CommandId(),
-                        model_card=ModelCard(
+                        model_meta=ModelMetadata(
                             model_id=ModelId("llama-3.2-1b"),
+                            pretty_name="Llama 3.2 1B",
                             n_layers=16,
                             storage_size=Memory.from_bytes(678948),
                             hidden_size=7168,
@@ -133,13 +143,9 @@ async def test_master():
                 command=(
                     ChatCompletion(
                         command_id=CommandId(),
-                        request_params=ChatCompletionTaskParams(
+                        request_params=ResponsesRequest(
                             model="llama-3.2-1b",
-                            messages=[
-                                ChatCompletionMessage(
-                                    role="user", content="Hello, how are you?"
-                                )
-                            ],
+                            input="Hello, how are you?",
                         ),
                     )
                 ),
@@ -153,7 +159,7 @@ async def test_master():
         assert events[0].idx == 0
         assert events[1].idx == 1
         assert events[2].idx == 2
-        assert isinstance(events[0].event, NodeGatheredInfo)
+        assert isinstance(events[0].event, NodePerformanceMeasured)
         assert isinstance(events[1].event, InstanceCreated)
         created_instance = events[1].event.instance
         assert isinstance(created_instance, MlxRingInstance)
@@ -166,8 +172,9 @@ async def test_master():
                     start_layer=0,
                     end_layer=16,
                     n_layers=16,
-                    model_card=ModelCard(
+                    model_meta=ModelMetadata(
                         model_id=ModelId("llama-3.2-1b"),
+                        pretty_name="Llama 3.2 1B",
                         n_layers=16,
                         storage_size=Memory.from_bytes(678948),
                         hidden_size=7168,
@@ -189,11 +196,9 @@ async def test_master():
         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(
+        assert events[2].event.task.task_params == ResponsesRequest(
             model="llama-3.2-1b",
-            messages=[
-                ChatCompletionMessage(role="user", content="Hello, how are you?")
-            ],
+            input="Hello, how are you?",
         )
 
         await master.shutdown()

src/exo/master/tests/test_openai_responses_api.py

@@ -0,0 +1,293 @@
+"""Tests for OpenAI Responses API types.
+
+ResponsesRequest is the canonical internal type used throughout the pipeline.
+No conversion is needed for Responses API requests.
+"""
+
+import json
+from typing import Any, cast
+
+import pydantic
+import pytest
+
+from exo.shared.types.openai_responses import (
+    ResponseCompletedEvent,
+    ResponseContentPartAddedEvent,
+    ResponseCreatedEvent,
+    ResponseInputMessage,
+    ResponseMessageItem,
+    ResponseOutputItemAddedEvent,
+    ResponseOutputItemDoneEvent,
+    ResponseOutputText,
+    ResponsesRequest,
+    ResponsesResponse,
+    ResponseTextDeltaEvent,
+    ResponseTextDoneEvent,
+    ResponseUsage,
+)
+
+
+class TestResponsesRequestAsCanonicalType:
+    """Tests for ResponsesRequest as the canonical internal type."""
+
+    def test_string_input(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input="Hello, how are you?",
+        )
+
+        assert request.model == "gpt-4o"
+        assert request.input == "Hello, how are you?"
+        assert request.instructions is None
+
+    def test_message_array_input(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input=[
+                ResponseInputMessage(role="user", content="Hello"),
+                ResponseInputMessage(role="assistant", content="Hi there!"),
+                ResponseInputMessage(role="user", content="How are you?"),
+            ],
+        )
+
+        assert isinstance(request.input, list)
+        assert len(request.input) == 3
+        assert request.input[0].role == "user"
+        assert request.input[0].content == "Hello"
+        assert request.input[1].role == "assistant"
+        assert request.input[1].content == "Hi there!"
+        assert request.input[2].role == "user"
+        assert request.input[2].content == "How are you?"
+
+    def test_request_with_instructions(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input="Hello",
+            instructions="You are a helpful assistant. Be concise.",
+        )
+
+        assert request.input == "Hello"
+        assert request.instructions == "You are a helpful assistant. Be concise."
+
+    def test_request_with_optional_parameters(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input="Hello",
+            max_output_tokens=500,
+            temperature=0.8,
+            top_p=0.95,
+            stream=True,
+        )
+
+        assert request.max_output_tokens == 500
+        assert request.temperature == 0.8
+        assert request.top_p == 0.95
+        assert request.stream is True
+
+    def test_request_with_new_fields(self):
+        """Test the additional fields added for internal use."""
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input="Hello",
+            top_k=40,
+            seed=42,
+            stop=["STOP", "END"],
+            tools=[{"type": "function", "function": {"name": "test"}}],
+        )
+
+        assert request.top_k == 40
+        assert request.seed == 42
+        assert request.stop == ["STOP", "END"]
+        assert request.tools == [{"type": "function", "function": {"name": "test"}}]
+
+    def test_request_with_system_role_in_messages(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input=[
+                ResponseInputMessage(role="system", content="Be helpful"),
+                ResponseInputMessage(role="user", content="Hello"),
+            ],
+        )
+
+        assert isinstance(request.input, list)
+        assert len(request.input) == 2
+        assert request.input[0].role == "system"
+        assert request.input[1].role == "user"
+
+    def test_request_with_developer_role(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input=[
+                ResponseInputMessage(role="developer", content="Internal note"),
+                ResponseInputMessage(role="user", content="Hello"),
+            ],
+        )
+
+        assert isinstance(request.input, list)
+        assert len(request.input) == 2
+        assert request.input[0].role == "developer"
+
+
+class TestResponsesRequestValidation:
+    """Tests for OpenAI Responses API request validation."""
+
+    def test_request_requires_model(self):
+        with pytest.raises(pydantic.ValidationError):
+            ResponsesRequest.model_validate(
+                {
+                    "input": "Hello",
+                }
+            )
+
+    def test_request_requires_input(self):
+        with pytest.raises(pydantic.ValidationError):
+            ResponsesRequest.model_validate(
+                {
+                    "model": "gpt-4o",
+                }
+            )
+
+    def test_request_accepts_string_input(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input="Hello",
+        )
+        assert request.input == "Hello"
+
+    def test_request_accepts_message_array_input(self):
+        request = ResponsesRequest(
+            model="gpt-4o",
+            input=[ResponseInputMessage(role="user", content="Hello")],
+        )
+        assert len(request.input) == 1
+
+
+class TestResponsesStreamingEvents:
+    """Tests for OpenAI Responses API streaming event serialization."""
+
+    def test_response_created_event_format(self):
+        response = ResponsesResponse(
+            id="resp_123",
+            model="gpt-4o",
+            status="in_progress",
+            output=[],
+            output_text="",
+        )
+        event = ResponseCreatedEvent(response=response)
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.created"
+        assert parsed["response"]["id"] == "resp_123"
+        assert parsed["response"]["object"] == "response"
+        assert parsed["response"]["status"] == "in_progress"
+
+    def test_output_item_added_event_format(self):
+        item = ResponseMessageItem(
+            id="item_123",
+            content=[ResponseOutputText(text="")],
+            status="in_progress",
+        )
+        event = ResponseOutputItemAddedEvent(output_index=0, item=item)
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.output_item.added"
+        assert parsed["output_index"] == 0
+        assert parsed["item"]["type"] == "message"
+        assert parsed["item"]["id"] == "item_123"
+        assert parsed["item"]["role"] == "assistant"
+
+    def test_content_part_added_event_format(self):
+        part = ResponseOutputText(text="")
+        event = ResponseContentPartAddedEvent(
+            output_index=0,
+            content_index=0,
+            part=part,
+        )
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.content_part.added"
+        assert parsed["output_index"] == 0
+        assert parsed["content_index"] == 0
+        assert parsed["part"]["type"] == "output_text"
+
+    def test_text_delta_event_format(self):
+        event = ResponseTextDeltaEvent(
+            output_index=0,
+            content_index=0,
+            delta="Hello",
+        )
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.output_text.delta"
+        assert parsed["output_index"] == 0
+        assert parsed["content_index"] == 0
+        assert parsed["delta"] == "Hello"
+
+    def test_text_done_event_format(self):
+        event = ResponseTextDoneEvent(
+            output_index=0,
+            content_index=0,
+            text="Hello, world!",
+        )
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.output_text.done"
+        assert parsed["text"] == "Hello, world!"
+
+    def test_output_item_done_event_format(self):
+        item = ResponseMessageItem(
+            id="item_123",
+            content=[ResponseOutputText(text="Hello, world!")],
+            status="completed",
+        )
+        event = ResponseOutputItemDoneEvent(output_index=0, item=item)
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.output_item.done"
+        assert parsed["item"]["status"] == "completed"
+        assert parsed["item"]["content"][0]["text"] == "Hello, world!"
+
+    def test_response_completed_event_format(self):
+        item = ResponseMessageItem(
+            id="item_123",
+            content=[ResponseOutputText(text="Hello!")],
+            status="completed",
+        )
+        response = ResponsesResponse(
+            id="resp_123",
+            model="gpt-4o",
+            status="completed",
+            output=[item],
+            output_text="Hello!",
+            usage=ResponseUsage(input_tokens=10, output_tokens=5, total_tokens=15),
+        )
+        event = ResponseCompletedEvent(response=response)
+        json_str = event.model_dump_json()
+        parsed = cast(dict[str, Any], json.loads(json_str))
+
+        assert parsed["type"] == "response.completed"
+        assert parsed["response"]["status"] == "completed"
+        assert parsed["response"]["output_text"] == "Hello!"
+        assert parsed["response"]["usage"]["total_tokens"] == 15
+
+    def test_sse_format(self):
+        """Test that SSE format is correctly generated."""
+        event = ResponseTextDeltaEvent(
+            output_index=0,
+            content_index=0,
+            delta="Hello",
+        )
+        # Simulate the SSE format used in the streaming generator
+        sse_line = (
+            f"event: response.output_text.delta\ndata: {event.model_dump_json()}\n\n"
+        )
+
+        assert sse_line.startswith("event: response.output_text.delta\n")
+        assert "data: " in sse_line
+        assert sse_line.endswith("\n\n")

src/exo/master/tests/test_placement.py

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

src/exo/master/tests/test_placement_utils.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 import pytest
 
 from exo.master.placement_utils import (
@@ -8,175 +10,154 @@ from exo.master.placement_utils import (
     get_shard_assignments,
     get_smallest_cycles,
 )
-from exo.master.tests.conftest import (
-    create_node_memory,
-    create_socket_connection,
-)
-from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import (
-    NetworkInterfaceInfo,
-    NodeNetworkInfo,
-)
-from exo.shared.types.topology import Connection, SocketConnection
+from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.profiling import NetworkInterfaceInfo, NodePerformanceProfile
+from exo.shared.types.topology import Connection, NodeInfo
 from exo.shared.types.worker.shards import Sharding
 
 
-def test_filter_cycles_by_memory():
+@pytest.fixture
+def topology() -> Topology:
+    topology = Topology()
+    return topology
+
+
+def test_filter_cycles_by_memory(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId], Connection],
+):
     # arrange
     node1_id = NodeId()
     node2_id = NodeId()
-    connection1 = Connection(
-        source=node1_id, sink=node2_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node2_id, sink=node1_id, edge=create_socket_connection(2)
-    )
 
-    node1_mem = create_node_memory(1000 * 1024)
-    node2_mem = create_node_memory(1000 * 1024)
-    node_memory = {node1_id: node1_mem, node2_id: node2_mem}
+    node1 = create_node(1000 * 1024, node1_id)
+    node2 = create_node(1000 * 1024, node2_id)
+
+    topology.add_node(node1)
+    topology.add_node(node2)
+
+    connection1 = create_connection(node1_id, node2_id)
+    connection2 = create_connection(node2_id, node1_id)
 
-    topology = Topology()
-    topology.add_node(node1_id)
-    topology.add_node(node2_id)
     topology.add_connection(connection1)
     topology.add_connection(connection2)
 
-    cycles = [c for c in topology.get_cycles() if len(c) != 1]
+    cycles = topology.get_cycles()
     assert len(cycles) == 1
     assert len(cycles[0]) == 2
 
     # act
-    filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_bytes(1))
+    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_bytes(1))
 
     # assert
     assert len(filtered_cycles) == 1
     assert len(filtered_cycles[0]) == 2
-    assert set(n for n in filtered_cycles[0]) == {node1_id, node2_id}
+    assert set(n.node_id for n in filtered_cycles[0]) == {node1_id, node2_id}
 
 
-def test_filter_cycles_by_insufficient_memory():
+def test_filter_cycles_by_insufficient_memory(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId], Connection],
+):
     # arrange
     node1_id = NodeId()
     node2_id = NodeId()
-    connection1 = Connection(
-        source=node1_id, sink=node2_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node2_id, sink=node1_id, edge=create_socket_connection(2)
-    )
 
-    node1_mem = create_node_memory(1000 * 1024)
-    node2_mem = create_node_memory(1000 * 1024)
-    node_memory = {node1_id: node1_mem, node2_id: node2_mem}
+    node1 = create_node(1000 * 1024, node1_id)
+    node2 = create_node(1000 * 1024, node2_id)
+
+    topology.add_node(node1)
+    topology.add_node(node2)
+
+    connection1 = create_connection(node1_id, node2_id)
+    connection2 = create_connection(node2_id, node1_id)
 
-    topology = Topology()
-    topology.add_node(node1_id)
-    topology.add_node(node2_id)
     topology.add_connection(connection1)
     topology.add_connection(connection2)
 
     # act
     filtered_cycles = filter_cycles_by_memory(
-        topology.get_cycles(), node_memory, Memory.from_kb(2001)
+        topology.get_cycles(), Memory.from_kb(2001)
     )
 
     # assert
     assert len(filtered_cycles) == 0
 
 
-def test_filter_multiple_cycles_by_memory():
+def test_filter_multiple_cycles_by_memory(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId], Connection],
+):
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
-    connection1 = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
-    )
-    connection3 = Connection(
-        source=node_a_id, sink=node_c_id, edge=create_socket_connection(3)
-    )
-    connection4 = Connection(
-        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
-    )
 
-    node_a_mem = create_node_memory(500 * 1024)
-    node_b_mem = create_node_memory(500 * 1024)
-    node_c_mem = create_node_memory(1000 * 1024)
-    node_memory = {
-        node_a_id: node_a_mem,
-        node_b_id: node_b_mem,
-        node_c_id: node_c_mem,
-    }
+    node_a = create_node(500 * 1024, node_a_id)
+    node_b = create_node(500 * 1024, node_b_id)
+    node_c = create_node(1000 * 1024, node_c_id)
 
-    topology = Topology()
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
-    topology.add_connection(connection3)
-    topology.add_connection(connection4)
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_node(node_c)
+
+    topology.add_connection(create_connection(node_a_id, node_b_id))
+    topology.add_connection(create_connection(node_b_id, node_a_id))
+
+    topology.add_connection(create_connection(node_a_id, node_c_id))
+    topology.add_connection(create_connection(node_c_id, node_b_id))
 
     cycles = topology.get_cycles()
 
     # act
-    filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_kb(1500))
+    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_kb(1500))
 
     # assert
     assert len(filtered_cycles) == 1
     assert len(filtered_cycles[0]) == 3
-    assert set(n for n in filtered_cycles[0]) == {
+    assert set(n.node_id for n in filtered_cycles[0]) == {
         node_a_id,
         node_b_id,
         node_c_id,
     }
 
 
-def test_get_smallest_cycles():
+def test_get_smallest_cycles(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId], Connection],
+):
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
 
-    topology = Topology()
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
+    node_a = create_node(500 * 1024, node_a_id)
+    node_b = create_node(500 * 1024, node_b_id)
+    node_c = create_node(1000 * 1024, node_c_id)
 
-    connection1 = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
-    )
-    connection3 = Connection(
-        source=node_a_id, sink=node_c_id, edge=create_socket_connection(3)
-    )
-    connection4 = Connection(
-        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
-    )
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_node(node_c)
 
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
-    topology.add_connection(connection3)
-    topology.add_connection(connection4)
-
-    cycles = [c for c in topology.get_cycles() if len(c) != 1]  # ignore singletons
+    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))
 
     # act
-    smallest_cycles = get_smallest_cycles(cycles)
+    smallest_cycles = get_smallest_cycles(topology.get_cycles())
 
     # assert
     assert len(smallest_cycles) == 1
     assert len(smallest_cycles[0]) == 2
-    assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id}
+    assert set(n.node_id for n in smallest_cycles[0]) == {node_a_id, node_b_id}
 
 
 @pytest.mark.parametrize(
@@ -191,6 +172,9 @@ def test_get_smallest_cycles():
     ],
 )
 def test_get_shard_assignments(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId], Connection],
     available_memory: tuple[int, int, int],
     total_layers: int,
     expected_layers: tuple[int, int, int],
@@ -200,61 +184,44 @@ def test_get_shard_assignments(
     node_b_id = NodeId()
     node_c_id = NodeId()
 
-    # create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists)
-    connection1 = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
-    )
-    connection3 = Connection(
-        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
-    )
-    connection4 = Connection(
-        source=node_b_id, sink=node_a_id, edge=create_socket_connection(4)
-    )
+    node_a = create_node(available_memory[0] * 1024, node_a_id)
+    node_b = create_node(available_memory[1] * 1024, node_b_id)
+    node_c = create_node(available_memory[2] * 1024, node_c_id)
 
-    topology = Topology()
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
-    topology.add_connection(connection3)
-    topology.add_connection(connection4)
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_node(node_c)
 
-    node_a_mem = create_node_memory(available_memory[0] * 1024)
-    node_b_mem = create_node_memory(available_memory[1] * 1024)
-    node_c_mem = create_node_memory(available_memory[2] * 1024)
-    node_memory = {
-        node_a_id: node_a_mem,
-        node_b_id: node_b_mem,
-        node_c_id: node_c_mem,
-    }
+    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_card = ModelCard(
+    model_meta = ModelMetadata(
         model_id=ModelId("test-model"),
+        pretty_name="Test Model",
         n_layers=total_layers,
         storage_size=Memory.from_kb(1000),
         hidden_size=1000,
         supports_tensor=True,
     )
-
     cycles = topology.get_cycles()
-
-    # pick the 3-node cycle deterministically (cycle ordering can vary)
-    selected_cycle = next(cycle for cycle in cycles if len(cycle) == 3)
+    selected_cycle = cycles[0]
 
     # act
     shard_assignments = get_shard_assignments(
-        model_card, selected_cycle, Sharding.Pipeline, node_memory=node_memory
+        model_meta, selected_cycle, Sharding.Pipeline
     )
 
     # assert
     runner_id_a = shard_assignments.node_to_runner[node_a_id]
     runner_id_b = shard_assignments.node_to_runner[node_b_id]
     runner_id_c = shard_assignments.node_to_runner[node_c_id]
-
+    assert (
+        shard_assignments.runner_to_shard[runner_id_c].end_layer
+        - shard_assignments.runner_to_shard[runner_id_c].start_layer
+        == expected_layers[2]
+    )
     assert (
         shard_assignments.runner_to_shard[runner_id_a].end_layer
         - shard_assignments.runner_to_shard[runner_id_a].start_layer
@@ -265,37 +232,30 @@ def test_get_shard_assignments(
         - shard_assignments.runner_to_shard[runner_id_b].start_layer
         == expected_layers[1]
     )
-    assert (
-        shard_assignments.runner_to_shard[runner_id_c].end_layer
-        - shard_assignments.runner_to_shard[runner_id_c].start_layer
-        == expected_layers[2]
-    )
 
 
-def test_get_hosts_from_subgraph():
+def test_get_hosts_from_subgraph(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId, int | None], Connection],
+):
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
-    topology = Topology()
 
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
+    node_a = create_node(500, node_a_id)
+    node_b = create_node(500, node_b_id)
+    node_c = create_node(1000, node_c_id)
 
-    connection1 = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
-    )
-    connection3 = Connection(
-        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
-    )
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_node(node_c)
 
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
-    topology.add_connection(connection3)
+    topology.add_connection(create_connection(node_a_id, node_b_id, 5001))
+    topology.add_connection(create_connection(node_b_id, node_c_id, 5002))
+    topology.add_connection(create_connection(node_c_id, node_a_id, 5003))
+    topology.add_connection(create_connection(node_b_id, node_a_id, 5004))
 
     # act
     hosts = get_hosts_from_subgraph(topology)
@@ -303,68 +263,95 @@ def test_get_hosts_from_subgraph():
     # assert
     assert len(hosts) == 3
     expected_hosts = [
-        Host(ip="169.254.0.1", port=1234),
-        Host(ip="169.254.0.2", port=1234),
-        Host(ip="169.254.0.3", port=1234),
+        Host(ip=("169.254.0.2"), port=5001),
+        Host(ip=("169.254.0.3"), port=5002),
+        Host(ip=("169.254.0.4"), port=5003),
     ]
     for expected_host in expected_hosts:
         assert expected_host in hosts
 
 
-def test_get_mlx_jaccl_coordinators():
+def test_get_mlx_jaccl_coordinators(
+    topology: Topology,
+    create_node: Callable[[int, NodeId | None], NodeInfo],
+    create_connection: Callable[[NodeId, NodeId, int | None], Connection],
+):
     # arrange
     node_a_id = NodeId()
     node_b_id = NodeId()
     node_c_id = NodeId()
 
-    # fully connected (directed) between the 3 nodes
-    conn_a_b = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    node_a = create_node(500 * 1024, node_a_id)
+    node_b = create_node(500 * 1024, node_b_id)
+    node_c = create_node(1000 * 1024, node_c_id)
+
+    conn_a_b = create_connection(node_a_id, node_b_id, 5001)
+    conn_b_a = create_connection(node_b_id, node_a_id, 5002)
+    conn_b_c = create_connection(node_b_id, node_c_id, 5003)
+    conn_c_b = create_connection(node_c_id, node_b_id, 5004)
+    conn_c_a = create_connection(node_c_id, node_a_id, 5005)
+    conn_a_c = create_connection(node_a_id, node_c_id, 5006)
+
+    # Update node profiles with network interfaces before adding to topology
+    assert node_a.node_profile is not None
+    assert node_b.node_profile is not None
+    assert node_c.node_profile is not None
+
+    node_a.node_profile = NodePerformanceProfile(
+        model_id="test",
+        chip_id="test",
+        friendly_name="test",
+        memory=node_a.node_profile.memory,
+        network_interfaces=[
+            NetworkInterfaceInfo(
+                name="en3",
+                ip_address=conn_a_b.send_back_multiaddr.ip_address,
+            ),
+            NetworkInterfaceInfo(
+                name="en4",
+                ip_address=conn_a_c.send_back_multiaddr.ip_address,
+            ),
+        ],
+        system=node_a.node_profile.system,
     )
-    conn_b_a = Connection(
-        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
+    node_b.node_profile = NodePerformanceProfile(
+        model_id="test",
+        chip_id="test",
+        friendly_name="test",
+        memory=node_b.node_profile.memory,
+        network_interfaces=[
+            NetworkInterfaceInfo(
+                name="en3",
+                ip_address=conn_b_a.send_back_multiaddr.ip_address,
+            ),
+            NetworkInterfaceInfo(
+                name="en4",
+                ip_address=conn_b_c.send_back_multiaddr.ip_address,
+            ),
+        ],
+        system=node_b.node_profile.system,
     )
-    conn_b_c = Connection(
-        source=node_b_id, sink=node_c_id, edge=create_socket_connection(3)
-    )
-    conn_c_b = Connection(
-        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
-    )
-    conn_c_a = Connection(
-        source=node_c_id, sink=node_a_id, edge=create_socket_connection(5)
-    )
-    conn_a_c = Connection(
-        source=node_a_id, sink=node_c_id, edge=create_socket_connection(6)
+    node_c.node_profile = NodePerformanceProfile(
+        model_id="test",
+        chip_id="test",
+        friendly_name="test",
+        memory=node_c.node_profile.memory,
+        network_interfaces=[
+            NetworkInterfaceInfo(
+                name="en3",
+                ip_address=conn_c_b.send_back_multiaddr.ip_address,
+            ),
+            NetworkInterfaceInfo(
+                name="en4",
+                ip_address=conn_c_a.send_back_multiaddr.ip_address,
+            ),
+        ],
+        system=node_c.node_profile.system,
     )
 
-    network_a = NodeNetworkInfo(
-        interfaces=[
-            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.5"),
-            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.2"),
-        ]
-    )
-    network_b = NodeNetworkInfo(
-        interfaces=[
-            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.1"),
-            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.4"),
-        ]
-    )
-    network_c = NodeNetworkInfo(
-        interfaces=[
-            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.3"),
-            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.6"),
-        ]
-    )
-    node_network = {
-        node_a_id: network_a,
-        node_b_id: network_b,
-        node_c_id: network_c,
-    }
-
-    topology = Topology()
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_node(node_c)
 
     topology.add_connection(conn_a_b)
     topology.add_connection(conn_b_a)
@@ -373,12 +360,11 @@ def test_get_mlx_jaccl_coordinators():
     topology.add_connection(conn_c_a)
     topology.add_connection(conn_a_c)
 
+    cycle = [node_a, node_b, node_c]
+
     # act
     coordinators = get_mlx_jaccl_coordinators(
-        node_a_id,
-        coordinator_port=5000,
-        cycle_digraph=topology,
-        node_network=node_network,
+        cycle, coordinator_port=5000, cycle_digraph=topology
     )
 
     # assert
@@ -399,22 +385,21 @@ def test_get_mlx_jaccl_coordinators():
             f"Coordinator for {node_id} should use port 5000"
         )
 
-    # Rank 0 (node_a) treats this as the listen socket so should listen on all IPs
+    # Rank 0 (node_a) treats this as the listen socket so should listen on all
+    # IPs
     assert coordinators[node_a_id].startswith("0.0.0.0:"), (
-        "Rank 0 node should use 0.0.0.0 as coordinator listen address"
+        "Rank 0 node should use localhost as coordinator"
     )
 
     # Non-rank-0 nodes should use the specific IP from their connection to rank 0
     # node_b uses the IP from conn_b_a (node_b -> node_a)
-    assert isinstance(conn_b_a.edge, SocketConnection)
-    assert (
-        coordinators[node_b_id] == f"{conn_b_a.edge.sink_multiaddr.ip_address}:5000"
+    assert coordinators[node_b_id] == (
+        f"{conn_b_a.send_back_multiaddr.ip_address}:5000"
     ), "node_b should use the IP from conn_b_a"
 
     # node_c uses the IP from conn_c_a (node_c -> node_a)
-    assert isinstance(conn_c_a.edge, SocketConnection)
     assert coordinators[node_c_id] == (
-        f"{conn_c_a.edge.sink_multiaddr.ip_address}:5000"
+        f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
     ), "node_c should use the IP from conn_c_a"
 
 
@@ -472,47 +457,33 @@ class TestAllocateLayersProportionally:
         assert sum(result) == 3
 
 
-def test_get_shard_assignments_insufficient_memory_raises():
+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()
-    topology = Topology()
 
     # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers)
-    node_a_mem = create_node_memory(900 * 1024)
-    node_b_mem = create_node_memory(50 * 1024)
-    node_c_mem = create_node_memory(10 * 1024)  # Insufficient memory
+    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_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_node(node_c)
 
-    conn_a_b = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    conn_b_c = Connection(
-        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
-    )
-    conn_c_a = Connection(
-        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
-    )
-    conn_b_a = Connection(
-        source=node_b_id, sink=node_a_id, edge=create_socket_connection(3)
-    )
-    topology.add_connection(conn_a_b)
-    topology.add_connection(conn_b_c)
-    topology.add_connection(conn_c_a)
-    topology.add_connection(conn_b_a)
+    topology.add_connection(create_connection(node_a_id, node_b_id))
+    topology.add_connection(create_connection(node_b_id, node_c_id))
+    topology.add_connection(create_connection(node_c_id, node_a_id))
+    topology.add_connection(create_connection(node_b_id, node_a_id))
 
-    node_memory = {
-        node_a_id: node_a_mem,
-        node_b_id: node_b_mem,
-        node_c_id: node_c_mem,
-    }
-
-    model_card = ModelCard(
+    model_meta = ModelMetadata(
         model_id=ModelId("test-model"),
+        pretty_name="Test Model",
         n_layers=20,
         storage_size=Memory.from_kb(1000),
         hidden_size=1000,
@@ -522,6 +493,4 @@ def test_get_shard_assignments_insufficient_memory_raises():
     selected_cycle = cycles[0]
 
     with pytest.raises(ValueError, match="insufficient memory"):
-        get_shard_assignments(
-            model_card, selected_cycle, Sharding.Pipeline, node_memory
-        )
+        get_shard_assignments(model_meta, selected_cycle, Sharding.Pipeline)

src/exo/master/tests/test_topology.py

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

src/exo/shared/apply.py

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

src/exo/shared/constants.py

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

src/exo/shared/logging.py

@@ -11,6 +11,9 @@ class InterceptLogger(HypercornLogger):
     def __init__(self, config: Config):
         super().__init__(config)
         assert self.error_logger
+        # TODO: Decide if we want to provide access logs
+        # assert self.access_logger
+        # self.access_logger.handlers = [_InterceptHandler()]
         self.error_logger.handlers = [_InterceptHandler()]
 
 

src/exo/shared/models/model_cards.py

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

src/exo/shared/models/model_meta.py

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

src/exo/shared/tests/conftest.py

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

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

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

src/exo/shared/tests/test_state_serialization.py

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

src/exo/shared/topology.py

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

src/exo/shared/types/api.py

@@ -4,9 +4,9 @@ from typing import Any, Literal
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core import PydanticUseDefault
 
-from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId
 from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 
@@ -157,10 +157,13 @@ class ChatCompletionTaskParams(BaseModel):
     stream: bool = False
     temperature: float | None = None
     top_p: float | None = None
+    top_k: int | None = None
     tools: list[dict[str, Any]] | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
+    # When True, continue the last assistant message without EOS tokens
+    continue_from_prefix: bool = False
 
 
 class BenchChatCompletionTaskParams(ChatCompletionTaskParams):
@@ -206,7 +209,7 @@ class DeleteInstanceTaskParams(BaseModel):
 class CreateInstanceResponse(BaseModel):
     message: str
     command_id: CommandId
-    model_card: ModelCard
+    model_meta: ModelMetadata
 
 
 class DeleteInstanceResponse(BaseModel):

src/exo/shared/types/chunks.py

@@ -1,10 +1,10 @@
 from enum import Enum
 
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import GenerationStats
+from exo.shared.types.api import GenerationStats, TopLogprobItem
 from exo.utils.pydantic_ext import TaggedModel
 
 from .api import FinishReason
+from .models import ModelId
 
 
 class ChunkType(str, Enum):
@@ -20,6 +20,8 @@ class BaseChunk(TaggedModel):
 class TokenChunk(BaseChunk):
     text: str
     token_id: int
+    logprob: float | None = None  # Log probability of the selected token
+    top_logprobs: list[TopLogprobItem] | None = None  # Top-k alternative tokens
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
     error_message: str | None = None

src/exo/shared/types/claude_api.py

@@ -0,0 +1,168 @@
+"""Claude Messages API types for request/response conversion."""
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+# Type aliases
+ClaudeRole = Literal["user", "assistant"]
+ClaudeStopReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]
+
+
+# Content block types
+class ClaudeTextBlock(BaseModel, frozen=True):
+    """Text content block in Claude Messages API."""
+
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ClaudeImageSource(BaseModel, frozen=True):
+    """Image source for Claude image blocks."""
+
+    type: Literal["base64", "url"]
+    media_type: str | None = None
+    data: str | None = None
+    url: str | None = None
+
+
+class ClaudeImageBlock(BaseModel, frozen=True):
+    """Image content block in Claude Messages API."""
+
+    type: Literal["image"] = "image"
+    source: ClaudeImageSource
+
+
+ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock
+
+
+# Request types
+class ClaudeMessage(BaseModel, frozen=True):
+    """Message in Claude Messages API request."""
+
+    role: ClaudeRole
+    content: str | list[ClaudeContentBlock]
+
+
+class ClaudeMessagesRequest(BaseModel):
+    """Request body for Claude Messages API."""
+
+    model: str
+    max_tokens: int
+    messages: list[ClaudeMessage]
+    system: str | list[ClaudeTextBlock] | None = None
+    stop_sequences: list[str] | None = None
+    stream: bool = False
+    temperature: float | None = None
+    top_p: float | None = None
+    top_k: int | None = None
+    metadata: dict[str, str] | None = None
+
+
+# Response types
+class ClaudeUsage(BaseModel, frozen=True):
+    """Token usage in Claude Messages API response."""
+
+    input_tokens: int
+    output_tokens: int
+
+
+class ClaudeMessagesResponse(BaseModel, frozen=True):
+    """Response body for Claude Messages API."""
+
+    id: str
+    type: Literal["message"] = "message"
+    role: Literal["assistant"] = "assistant"
+    content: list[ClaudeTextBlock]
+    model: str
+    stop_reason: ClaudeStopReason | None = None
+    stop_sequence: str | None = None
+    usage: ClaudeUsage
+
+
+# Streaming event types
+class ClaudeMessageStart(BaseModel, frozen=True):
+    """Partial message in message_start event."""
+
+    id: str
+    type: Literal["message"] = "message"
+    role: Literal["assistant"] = "assistant"
+    content: list[ClaudeTextBlock] = Field(default_factory=list)
+    model: str
+    stop_reason: ClaudeStopReason | None = None
+    stop_sequence: str | None = None
+    usage: ClaudeUsage
+
+
+class ClaudeMessageStartEvent(BaseModel, frozen=True):
+    """Event sent at start of message stream."""
+
+    type: Literal["message_start"] = "message_start"
+    message: ClaudeMessageStart
+
+
+class ClaudeContentBlockStartEvent(BaseModel, frozen=True):
+    """Event sent at start of a content block."""
+
+    type: Literal["content_block_start"] = "content_block_start"
+    index: int
+    content_block: ClaudeTextBlock
+
+
+class ClaudeTextDelta(BaseModel, frozen=True):
+    """Delta for text content block."""
+
+    type: Literal["text_delta"] = "text_delta"
+    text: str
+
+
+class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True):
+    """Event sent for content block delta."""
+
+    type: Literal["content_block_delta"] = "content_block_delta"
+    index: int
+    delta: ClaudeTextDelta
+
+
+class ClaudeContentBlockStopEvent(BaseModel, frozen=True):
+    """Event sent at end of a content block."""
+
+    type: Literal["content_block_stop"] = "content_block_stop"
+    index: int
+
+
+class ClaudeMessageDeltaUsage(BaseModel, frozen=True):
+    """Usage in message_delta event."""
+
+    output_tokens: int
+
+
+class ClaudeMessageDelta(BaseModel, frozen=True):
+    """Delta in message_delta event."""
+
+    stop_reason: ClaudeStopReason | None = None
+    stop_sequence: str | None = None
+
+
+class ClaudeMessageDeltaEvent(BaseModel, frozen=True):
+    """Event sent with final message delta."""
+
+    type: Literal["message_delta"] = "message_delta"
+    delta: ClaudeMessageDelta
+    usage: ClaudeMessageDeltaUsage
+
+
+class ClaudeMessageStopEvent(BaseModel, frozen=True):
+    """Event sent at end of message stream."""
+
+    type: Literal["message_stop"] = "message_stop"
+
+
+ClaudeStreamEvent = (
+    ClaudeMessageStartEvent
+    | ClaudeContentBlockStartEvent
+    | ClaudeContentBlockDeltaEvent
+    | ClaudeContentBlockStopEvent
+    | ClaudeMessageDeltaEvent
+    | ClaudeMessageStopEvent
+)

src/exo/shared/types/commands.py

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

src/exo/shared/types/common.py

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

src/exo/shared/types/events.py

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

src/exo/shared/types/models.py

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

src/exo/shared/types/multiaddr.py

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

src/exo/shared/types/openai_responses.py

@@ -0,0 +1,190 @@
+"""OpenAI Responses API types for request/response conversion.
+
+ResponsesRequest serves as both:
+1. The external API request type for /v1/responses
+2. The canonical internal type used throughout the inference pipeline
+
+All external API formats (Chat Completions, Claude) are converted to
+ResponsesRequest at the API boundary.
+"""
+
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+# Type aliases
+ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
+ResponseRole = Literal["user", "assistant", "system", "developer"]
+
+
+# Request types
+class ResponseInputMessage(BaseModel, frozen=True):
+    """Input message for Responses API.
+
+    This is also used as the internal message format throughout the pipeline.
+    """
+
+    role: ResponseRole
+    content: str
+
+
+class ResponsesRequest(BaseModel):
+    """Request body for OpenAI Responses API.
+
+    This is also the canonical internal task params format used throughout
+    the inference pipeline. All external API formats are converted to this
+    format at the API boundary.
+
+    Field mapping from other APIs:
+    - input: Replaces 'messages' from Chat Completions
+    - instructions: System message, extracted from messages or Claude's 'system'
+    - max_output_tokens: Replaces 'max_tokens' from Chat Completions
+    """
+
+    model: str
+    input: str | list[ResponseInputMessage]
+    instructions: str | None = None
+    max_output_tokens: int | None = None
+    temperature: float | None = None
+    top_p: float | None = None
+    top_k: int | None = None
+    stop: str | list[str] | None = None
+    seed: int | None = None
+    stream: bool = False
+    # Tools support
+    tools: list[dict[str, Any]] | None = None
+    # previous_response_id not supported in MVP
+    metadata: dict[str, str] | None = None
+    # When True, continue the last assistant message without EOS tokens
+    continue_from_prefix: bool = False
+
+
+# Response types
+class ResponseOutputText(BaseModel, frozen=True):
+    """Text content in response output."""
+
+    type: Literal["output_text"] = "output_text"
+    text: str
+    annotations: list[dict[str, str]] = Field(default_factory=list)
+
+
+class ResponseMessageItem(BaseModel, frozen=True):
+    """Message item in response output array."""
+
+    type: Literal["message"] = "message"
+    id: str
+    role: Literal["assistant"] = "assistant"
+    content: list[ResponseOutputText]
+    status: ResponseStatus = "completed"
+
+
+ResponseItem = ResponseMessageItem  # Can expand for function_call, reasoning, etc.
+
+
+class ResponseUsage(BaseModel, frozen=True):
+    """Token usage in Responses API response."""
+
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+
+
+class ResponsesResponse(BaseModel, frozen=True):
+    """Response body for OpenAI Responses API."""
+
+    id: str
+    object: Literal["response"] = "response"
+    created_at: int = Field(default_factory=lambda: int(time.time()))
+    status: ResponseStatus = "completed"
+    model: str
+    output: list[ResponseItem]
+    output_text: str
+    usage: ResponseUsage | None = None
+
+
+# Streaming event types
+class ResponseCreatedEvent(BaseModel, frozen=True):
+    """Event sent when response is created."""
+
+    type: Literal["response.created"] = "response.created"
+    response: ResponsesResponse
+
+
+class ResponseInProgressEvent(BaseModel, frozen=True):
+    """Event sent when response starts processing."""
+
+    type: Literal["response.in_progress"] = "response.in_progress"
+    response: ResponsesResponse
+
+
+class ResponseOutputItemAddedEvent(BaseModel, frozen=True):
+    """Event sent when an output item is added."""
+
+    type: Literal["response.output_item.added"] = "response.output_item.added"
+    output_index: int
+    item: ResponseItem
+
+
+class ResponseContentPartAddedEvent(BaseModel, frozen=True):
+    """Event sent when a content part is added."""
+
+    type: Literal["response.content_part.added"] = "response.content_part.added"
+    output_index: int
+    content_index: int
+    part: ResponseOutputText
+
+
+class ResponseTextDeltaEvent(BaseModel, frozen=True):
+    """Event sent for text delta during streaming."""
+
+    type: Literal["response.output_text.delta"] = "response.output_text.delta"
+    output_index: int
+    content_index: int
+    delta: str
+
+
+class ResponseTextDoneEvent(BaseModel, frozen=True):
+    """Event sent when text content is done."""
+
+    type: Literal["response.output_text.done"] = "response.output_text.done"
+    output_index: int
+    content_index: int
+    text: str
+
+
+class ResponseContentPartDoneEvent(BaseModel, frozen=True):
+    """Event sent when a content part is done."""
+
+    type: Literal["response.content_part.done"] = "response.content_part.done"
+    output_index: int
+    content_index: int
+    part: ResponseOutputText
+
+
+class ResponseOutputItemDoneEvent(BaseModel, frozen=True):
+    """Event sent when an output item is done."""
+
+    type: Literal["response.output_item.done"] = "response.output_item.done"
+    output_index: int
+    item: ResponseItem
+
+
+class ResponseCompletedEvent(BaseModel, frozen=True):
+    """Event sent when response is completed."""
+
+    type: Literal["response.completed"] = "response.completed"
+    response: ResponsesResponse
+
+
+ResponsesStreamEvent = (
+    ResponseCreatedEvent
+    | ResponseInProgressEvent
+    | ResponseOutputItemAddedEvent
+    | ResponseContentPartAddedEvent
+    | ResponseTextDeltaEvent
+    | ResponseTextDoneEvent
+    | ResponseContentPartDoneEvent
+    | ResponseOutputItemDoneEvent
+    | ResponseCompletedEvent
+)

src/exo/shared/types/profiling.py

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

src/exo/shared/types/state.py

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

src/exo/shared/types/tasks.py

@@ -2,8 +2,8 @@ from enum import Enum
 
 from pydantic import Field
 
-from exo.shared.types.api import ChatCompletionTaskParams
 from exo.shared.types.common import CommandId, Id
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.worker.instances import BoundInstance, InstanceId
 from exo.shared.types.worker.runners import RunnerId
 from exo.shared.types.worker.shards import ShardMetadata
@@ -50,7 +50,7 @@ class StartWarmup(BaseTask):  # emitted by Worker
 
 class ChatCompletion(BaseTask):  # emitted by Master
     command_id: CommandId
-    task_params: ChatCompletionTaskParams
+    task_params: ResponsesRequest
 
     error_type: str | None = Field(default=None)
     error_message: str | None = Field(default=None)

src/exo/shared/types/thunderbolt.py

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

src/exo/shared/types/topology.py

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

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

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

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

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

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

@@ -1,4 +1,4 @@
-from exo.shared.types.api import FinishReason, GenerationStats
+from exo.shared.types.api import FinishReason, GenerationStats, TopLogprobItem
 from exo.utils.pydantic_ext import TaggedModel
 
 
@@ -13,7 +13,8 @@ class TokenizedResponse(BaseRunnerResponse):
 class GenerationResponse(BaseRunnerResponse):
     text: str
     token: int
-    # logprobs: list[float] | None = None # too big. we can change to be top-k
+    logprob: float | None = None  # Log probability of the selected token
+    top_logprobs: list[TopLogprobItem] | None = None  # Top-k alternative tokens
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
 

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

@@ -2,8 +2,8 @@ from collections.abc import Mapping
 
 from pydantic import model_validator
 
-from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import Id, NodeId
+from exo.shared.types.models import ModelId
 from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 

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

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

src/exo/utils/info_gatherer/info_gatherer.py

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

src/exo/utils/info_gatherer/macmon.py

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

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

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

src/exo/utils/pydantic_ext.py

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

src/exo/utils/tests/testing_mp.py

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

src/exo/worker/download/download_utils.py

@@ -5,7 +5,6 @@ import shutil
 import ssl
 import time
 import traceback
-from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -460,10 +459,10 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
     # (iii) Tensor parallel requires all files.
     return ["*"]
     try:
-        weight_map = await get_weight_map(str(shard.model_card.model_id))
+        weight_map = await get_weight_map(str(shard.model_meta.model_id))
         return get_allow_patterns(weight_map, shard)
     except Exception:
-        logger.error(f"Error getting weight map for {shard.model_card.model_id=}")
+        logger.error(f"Error getting weight map for {shard.model_meta.model_id=}")
         logger.error(traceback.format_exc())
         return ["*"]
 
@@ -526,24 +525,24 @@ async def download_progress_for_local_path(
 
 async def download_shard(
     shard: ShardMetadata,
-    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
     max_parallel_downloads: int = 8,
     skip_download: bool = False,
     allow_patterns: list[str] | None = None,
 ) -> tuple[Path, RepoDownloadProgress]:
     if not skip_download:
-        logger.info(f"Downloading {shard.model_card.model_id=}")
+        logger.info(f"Downloading {shard.model_meta.model_id=}")
 
     # Handle local paths
-    if await aios.path.exists(str(shard.model_card.model_id)):
-        logger.info(f"Using local model path {shard.model_card.model_id}")
-        local_path = Path(str(shard.model_card.model_id))
+    if await aios.path.exists(str(shard.model_meta.model_id)):
+        logger.info(f"Using local model path {shard.model_meta.model_id}")
+        local_path = Path(str(shard.model_meta.model_id))
         return local_path, await download_progress_for_local_path(
-            str(shard.model_card.model_id), shard, local_path
+            str(shard.model_meta.model_id), shard, local_path
         )
 
     revision = "main"
-    target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
+    target_dir = await ensure_models_dir() / str(shard.model_meta.model_id).replace(
         "/", "--"
     )
     if not skip_download:
@@ -552,13 +551,13 @@ async def download_shard(
     if not allow_patterns:
         allow_patterns = await resolve_allow_patterns(shard)
 
-    logger.info(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
+    logger.info(f"Downloading {shard.model_meta.model_id=} with {allow_patterns=}")
 
     all_start_time = time.time()
     # TODO: currently not recursive. Some models might require subdirectories - thus this will need to be changed.
     #  Update: <- This does not seem to be the case. Yay?
     file_list = await fetch_file_list_with_cache(
-        str(shard.model_card.model_id), revision, recursive=True
+        str(shard.model_meta.model_id), revision, recursive=True
     )
     filtered_file_list = list(
         filter_repo_objects(
@@ -567,9 +566,9 @@ async def download_shard(
     )
     file_progress: dict[str, RepoFileDownloadProgress] = {}
 
-    async def on_progress_wrapper(
+    def on_progress_wrapper(
         file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ) -> None:
+    ):
         start_time = (
             file_progress[file.path].start_time
             if file.path in file_progress
@@ -592,7 +591,7 @@ async def download_shard(
             else timedelta(seconds=0)
         )
         file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=str(shard.model_card.model_id),
+            repo_id=str(shard.model_meta.model_id),
             repo_revision=revision,
             file_path=file.path,
             downloaded=Memory.from_bytes(curr_bytes),
@@ -605,11 +604,11 @@ async def download_shard(
             else "in_progress",
             start_time=start_time,
         )
-        await on_progress(
+        on_progress(
             shard,
             calculate_repo_progress(
                 shard,
-                str(shard.model_card.model_id),
+                str(shard.model_meta.model_id),
                 revision,
                 file_progress,
                 all_start_time,
@@ -619,7 +618,7 @@ async def download_shard(
     for file in filtered_file_list:
         downloaded_bytes = await get_downloaded_size(target_dir / file.path)
         file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=str(shard.model_card.model_id),
+            repo_id=str(shard.model_meta.model_id),
             repo_revision=revision,
             file_path=file.path,
             downloaded=Memory.from_bytes(downloaded_bytes),
@@ -633,21 +632,14 @@ async def download_shard(
 
     semaphore = asyncio.Semaphore(max_parallel_downloads)
 
-    def schedule_progress(
-        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ) -> None:
-        asyncio.create_task(
-            on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed)
-        )
-
-    async def download_with_semaphore(file: FileListEntry) -> None:
+    async def download_with_semaphore(file: FileListEntry):
         async with semaphore:
             await download_file_with_retry(
-                str(shard.model_card.model_id),
+                str(shard.model_meta.model_id),
                 revision,
                 file.path,
                 target_dir,
-                lambda curr_bytes, total_bytes, is_renamed: schedule_progress(
+                lambda curr_bytes, total_bytes, is_renamed: on_progress_wrapper(
                     file, curr_bytes, total_bytes, is_renamed
                 ),
             )
@@ -657,9 +649,9 @@ async def download_shard(
             *[download_with_semaphore(file) for file in filtered_file_list]
         )
     final_repo_progress = calculate_repo_progress(
-        shard, str(shard.model_card.model_id), revision, file_progress, all_start_time
+        shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
     )
-    await on_progress(shard, final_repo_progress)
+    on_progress(shard, final_repo_progress)
     if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
         return target_dir / gguf.path, final_repo_progress
     else:

src/exo/worker/download/impl_shard_downloader.py

@@ -1,10 +1,9 @@
 import asyncio
-from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
 from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.models.model_meta import get_model_card
+from exo.shared.models.model_meta import get_model_meta
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
@@ -20,21 +19,21 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
 
 
 async def build_base_shard(model_id: str) -> ShardMetadata:
-    model_card = await get_model_card(model_id)
+    model_meta = await get_model_meta(model_id)
     return PipelineShardMetadata(
-        model_card=model_card,
+        model_meta=model_meta,
         device_rank=0,
         world_size=1,
         start_layer=0,
-        end_layer=model_card.n_layers,
-        n_layers=model_card.n_layers,
+        end_layer=model_meta.n_layers,
+        n_layers=model_meta.n_layers,
     )
 
 
 async def build_full_shard(model_id: str) -> PipelineShardMetadata:
     base_shard = await build_base_shard(model_id)
     return PipelineShardMetadata(
-        model_card=base_shard.model_card,
+        model_meta=base_shard.model_meta,
         device_rank=base_shard.device_rank,
         world_size=base_shard.world_size,
         start_layer=base_shard.start_layer,
@@ -49,8 +48,7 @@ class SingletonShardDownloader(ShardDownloader):
         self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
 
     def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
     ) -> None:
         self.shard_downloader.on_progress(callback)
 
@@ -85,19 +83,18 @@ class CachedShardDownloader(ShardDownloader):
         self.cache: dict[tuple[str, ShardMetadata], Path] = {}
 
     def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
     ) -> None:
         self.shard_downloader.on_progress(callback)
 
     async def ensure_shard(
         self, shard: ShardMetadata, config_only: bool = False
     ) -> Path:
-        if (shard.model_card.model_id, shard) in self.cache:
-            return self.cache[(shard.model_card.model_id, shard)]
+        if (shard.model_meta.model_id, shard) in self.cache:
+            return self.cache[(shard.model_meta.model_id, shard)]
 
         target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
-        self.cache[(shard.model_card.model_id, shard)] = target_dir
+        self.cache[(shard.model_meta.model_id, shard)] = target_dir
         return target_dir
 
     async def get_shard_download_status(
@@ -116,18 +113,17 @@ class ResumableShardDownloader(ShardDownloader):
     def __init__(self, max_parallel_downloads: int = 8):
         self.max_parallel_downloads = max_parallel_downloads
         self.on_progress_callbacks: list[
-            Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
+            Callable[[ShardMetadata, RepoDownloadProgress], None]
         ] = []
 
-    async def on_progress_wrapper(
+    def on_progress_wrapper(
         self, shard: ShardMetadata, progress: RepoDownloadProgress
     ) -> None:
         for callback in self.on_progress_callbacks:
-            await callback(shard, progress)
+            callback(shard, progress)
 
     def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
     ) -> None:
         self.on_progress_callbacks.append(callback)
 

src/exo/worker/download/shard_downloader.py

@@ -1,12 +1,11 @@
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable
 from copy import copy
 from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
-from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
@@ -32,8 +31,7 @@ class ShardDownloader(ABC):
 
     @abstractmethod
     def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
     ) -> None:
         pass
 
@@ -61,8 +59,7 @@ class NoopShardDownloader(ShardDownloader):
         return Path("/tmp/noop_shard")
 
     def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
     ) -> None:
         pass
 
@@ -86,8 +83,9 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
     repo_id="noop",
     repo_revision="noop",
     shard=PipelineShardMetadata(
-        model_card=ModelCard(
+        model_meta=ModelMetadata(
             model_id=ModelId("noop"),
+            pretty_name="noope",
             storage_size=Memory.from_bytes(0),
             n_layers=1,
             hidden_size=1,

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

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

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

@@ -1,10 +1,7 @@
-import os
-import threading
 from abc import ABC, abstractmethod
-from collections.abc import Callable
 from functools import partial
 from inspect import signature
-from typing import TYPE_CHECKING, Any, Protocol, cast
+from typing import TYPE_CHECKING, Callable, Protocol, cast
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -32,40 +29,6 @@ from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
 from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
 
-TimeoutCallback = Callable[[], None]
-
-
-def eval_with_timeout(
-    mlx_item: Any,  # pyright: ignore[reportAny]
-    timeout_seconds: float = 60.0,
-    on_timeout: TimeoutCallback | None = None,
-) -> None:
-    """Evaluate MLX item with a hard timeout.
-
-    If on_timeout callback is provided, it will be called before terminating
-    the process. This allows the runner to send a failure event before exit.
-    """
-    completed = threading.Event()
-
-    def watchdog() -> None:
-        if not completed.wait(timeout=timeout_seconds):
-            logger.error(
-                f"mlx_item evaluation timed out after {timeout_seconds:.0f}s. "
-                "This may indicate an issue with FAST_SYNCH and tensor parallel sharding. "
-                "Terminating process."
-            )
-            if on_timeout is not None:
-                on_timeout()
-            os._exit(1)
-
-    watchdog_thread = threading.Thread(target=watchdog, daemon=True)
-    watchdog_thread.start()
-
-    try:
-        mx.eval(mlx_item)  # pyright: ignore[reportAny]
-    finally:
-        completed.set()
-
 
 class _LayerCallable(Protocol):
     """Structural type that any compatible layer must satisfy.
@@ -83,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:
@@ -97,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):
@@ -145,6 +106,7 @@ class PipelineLastLayer(CustomMlxLayer):
             if cache is not None:
                 cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
 
+        output = mx.distributed.all_gather(output, group=self.group)[-output.shape[0] :]
         return output
 
 
@@ -206,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)
@@ -229,70 +181,12 @@ def pipeline_auto_parallel(
         "Expected a list of layers after auto-parallel initialisation"
     )
 
-    return patch_pipeline_model(model, group)
-
-
-def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
-    # Patch __call__ on the model's class
-    cls = model.__class__
-    original_call = cls.__call__  # type :ignore
-    call_signature = signature(original_call)  # type :ignore
-
-    def patched_call(
-        self: T,
-        *args: object,
-        **kwargs: object,
-    ) -> mx.array:
-        logits: mx.array = original_call(self, *args, **kwargs)  # type: ignore
-        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
-            "cache", None
-        )
-
-        # Add dependency to last cache entry to ensure distributed ops are evaluated
-        if cache is not None:
-            cache[-1].state = mx.depends(cache[-1].state, logits)  # type: ignore
-
-        logits = mx.distributed.all_gather(logits, group=group)[
-            -logits.shape[0] :
-        ]  # type :ignore
-
-        return logits
-
-    cls.__call__ = patched_call
-    return model
-
-
-def patch_tensor_model[T](model: T) -> T:
-    """Patch model's __call__ to ensure distributed ops sync during inference."""
-    cls = model.__class__
-    original_call = cls.__call__
-    call_signature = signature(original_call)
-
-    def patched_call(
-        self: T,
-        *args: object,
-        **kwargs: object,
-    ) -> mx.array:
-        logits: mx.array = original_call(self, *args, **kwargs)  # pyright: ignore[reportAny]
-        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
-            "cache", None
-        )
-
-        # Add dependency to last cache entry to ensure distributed ops are evaluated
-        if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
-            cache[-1].state = mx.depends(cache[-1].state, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
-
-        return logits
-
-    cls.__call__ = patched_call
     return model
 
 
 def tensor_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
-    timeout_seconds: float = 60.0,
-    on_timeout: TimeoutCallback | None = None,
 ) -> nn.Module:
     all_to_sharded_linear = partial(
         shard_linear,
@@ -337,7 +231,7 @@ def tensor_auto_parallel(
     if hasattr(model, "shard"):
         try:
             model.shard(group)  # type: ignore
-            return patch_tensor_model(model)
+            return model
         except (AttributeError, TypeError, NameError):
             pass
 
@@ -387,10 +281,7 @@ def tensor_auto_parallel(
     else:
         raise ValueError(f"Unsupported model type: {type(model)}")
 
-    model = tensor_parallel_sharding_strategy.shard_model(
-        model, timeout_seconds, on_timeout
-    )
-    return patch_tensor_model(model)
+    return tensor_parallel_sharding_strategy.shard_model(model)
 
 
 class TensorParallelShardingStrategy(ABC):
@@ -410,27 +301,13 @@ class TensorParallelShardingStrategy(ABC):
         self.N = group.size()
 
     @abstractmethod
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module: ...
+    def shard_model(self, model: nn.Module) -> nn.Module: ...
 
 
 class LlamaShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module:
+    def shard_model(self, model: nn.Module) -> nn.Module:
         model = cast(LlamaModel, model)
         for layer in model.layers:
-            # Force load weights before sharding to avoid FAST_SYNCH deadlock
-            eval_with_timeout(
-                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
-            )
             layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
             layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -473,17 +350,9 @@ def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
 
 
 class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module:
+    def shard_model(self, model: nn.Module) -> nn.Module:
         model = cast(DeepseekV3Model, model)
         for layer in model.layers:
-            eval_with_timeout(
-                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
-            )
             # Shard the self attention
             if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
@@ -535,17 +404,9 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
 
 
 class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module:
+    def shard_model(self, model: nn.Module) -> nn.Module:
         model = cast(MiniMaxModel, model)
         for layer in model.layers:
-            eval_with_timeout(
-                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
-            )
             # Shard the self attention
             layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -566,23 +427,15 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
                 layer.block_sparse_moe.switch_mlp.up_proj
             )
             layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
-            layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.block_sparse_moe.sharding_group = self.group
 
         return model
 
 
 class QwenShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module:
+    def shard_model(self, model: nn.Module) -> nn.Module:
         model = cast(Qwen3MoeModel, model)
         for layer in model.layers:
-            eval_with_timeout(
-                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
-            )
             # Shard the self attention
             layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -626,18 +479,10 @@ class ShardedQwenMoE(CustomMlxLayer):
 
 
 class GptOssShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module:
+    def shard_model(self, model: nn.Module) -> nn.Module:
         model = cast(GptOssMoeModel, model)
 
         for layer in model.layers:
-            eval_with_timeout(
-                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
-            )
             layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
             layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -661,7 +506,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
 
             layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
-            layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.mlp.sharding_group = self.group
 
         return model
 

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

@@ -8,13 +8,12 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 # from exo.engines.mlx.cache import KVPrefixCache
 from exo.shared.types.api import (
-    BenchChatCompletionTaskParams,
-    ChatCompletionMessage,
     FinishReason,
     GenerationStats,
+    TopLogprobItem,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
@@ -53,14 +52,9 @@ def warmup_inference(
 
     warmup_prompt = apply_chat_template(
         tokenizer=tokenizer,
-        chat_task_data=ChatCompletionTaskParams(
+        task_params=ResponsesRequest(
             model="",
-            messages=[
-                ChatCompletionMessage(
-                    role="user",
-                    content=content,
-                )
-            ],
+            input=content,
         ),
     )
 
@@ -115,15 +109,68 @@ def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
     return eos
 
 
+def extract_top_logprobs(
+    logprobs: mx.array,
+    tokenizer: TokenizerWrapper,
+    top_k: int,
+    selected_token: int,
+) -> tuple[float, list[TopLogprobItem]]:
+    """Extract the selected token's logprob and top-k alternative tokens.
+
+    Args:
+        logprobs: Full vocabulary logprobs array from MLX
+        tokenizer: Tokenizer for decoding token IDs to strings
+        top_k: Number of top alternatives to return
+        selected_token: The token ID that was actually sampled
+
+    Returns:
+        Tuple of (selected_token_logprob, list of TopLogprobItem for top-k tokens)
+    """
+    # Get the logprob of the selected token
+    selected_logprob = float(logprobs[selected_token].item())
+
+    # Get top-k indices (most probable tokens)
+    # mx.argpartition gives indices that would partition the array
+    # We negate logprobs since argpartition finds smallest, and we want largest
+    top_k = min(top_k, logprobs.shape[0])  # Don't exceed vocab size
+    top_indices = mx.argpartition(-logprobs, top_k)[:top_k]
+
+    # Get the actual logprob values for these indices
+    top_values = logprobs[top_indices]
+
+    # Sort by logprob (descending) for consistent ordering
+    sort_order = mx.argsort(-top_values)
+    top_indices = top_indices[sort_order]
+    top_values = top_values[sort_order]
+
+    # Convert to list of TopLogprobItem
+    top_logprob_items: list[TopLogprobItem] = []
+    for i in range(top_k):
+        token_id = int(top_indices[i].item())
+        token_logprob = float(top_values[i].item())
+        # Decode token ID to string
+        token_str = tokenizer.decode([token_id])
+        # Get byte representation
+        token_bytes = list(token_str.encode("utf-8"))
+        top_logprob_items.append(
+            TopLogprobItem(
+                token=token_str,
+                logprob=token_logprob,
+                bytes=token_bytes,
+            )
+        )
+
+    return selected_logprob, top_logprob_items
+
+
 def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
-    task: ChatCompletionTaskParams,
-    prompt: str,
+    task: ResponsesRequest,
+    is_bench: bool = False,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
-    is_bench: bool = isinstance(task, BenchChatCompletionTaskParams)
 
     # Currently we support chat-completion tasks only.
     logger.info(f"task_params: {task}")
@@ -131,6 +178,11 @@ def mlx_generate(
     if task.seed is not None:
         mx.random.seed(task.seed)
 
+    prompt = apply_chat_template(
+        tokenizer=tokenizer,
+        task_params=task,
+    )
+
     caches = make_kv_cache(model=model)
 
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
@@ -142,9 +194,20 @@ def mlx_generate(
     sampler = make_sampler(
         temp=task.temperature if task.temperature is not None else 0.7,
         top_p=task.top_p if task.top_p is not None else 1.0,
+        top_k=task.top_k if task.top_k is not None else 0,
     )
 
-    max_tokens = task.max_tokens or MAX_TOKENS
+    # Normalize stop sequences to a list
+    stop_sequences: list[str] = (
+        ([task.stop] if isinstance(task.stop, str) else task.stop)
+        if task.stop is not None
+        else []
+    )
+    max_stop_len = max((len(s) for s in stop_sequences), default=0)
+
+    max_tokens = task.max_output_tokens or MAX_TOKENS
+    accumulated_text = ""
+
     for out in stream_generate(
         model=model,
         tokenizer=tokenizer,
@@ -153,15 +216,35 @@ def mlx_generate(
         sampler=sampler,
         logits_processors=logits_processors,
         prompt_cache=caches,
-        # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
         prefill_step_size=2048,
         kv_group_size=KV_GROUP_SIZE,
         kv_bits=KV_BITS,
     ):
         logger.info(out.text)
+        accumulated_text += out.text
 
+        # Check for stop sequences
+        text = out.text
+        finish_reason: FinishReason | None = cast(
+            FinishReason | None, out.finish_reason
+        )
+        stop_matched = False
+
+        if stop_sequences:
+            for stop_seq in stop_sequences:
+                if stop_seq in accumulated_text:
+                    # Trim text to just before the stop sequence
+                    stop_index = accumulated_text.find(stop_seq)
+                    text_before_stop = accumulated_text[:stop_index]
+                    chunk_start = len(accumulated_text) - len(out.text)
+                    text = text_before_stop[chunk_start:]
+                    finish_reason = "stop"
+                    stop_matched = True
+                    break
+
+        is_done = finish_reason is not None
         stats: GenerationStats | None = None
-        if out.finish_reason is not None:
+        if is_done:
             stats = GenerationStats(
                 prompt_tps=float(out.prompt_tps),
                 generation_tps=float(out.generation_tps),
@@ -169,22 +252,33 @@ def mlx_generate(
                 generation_tokens=int(out.generation_tokens),
                 peak_memory_usage=Memory.from_gb(out.peak_memory),
             )
-
-            if out.finish_reason not in get_args(FinishReason):
-                # We don't throw here as this failure case is really not all that bad
-                # Just log the error and move on
+            if not stop_matched and out.finish_reason not in get_args(FinishReason):
                 logger.warning(
                     f"Model generated unexpected finish_reason: {out.finish_reason}"
                 )
 
+        # Extract logprobs from the full vocabulary logprobs array
+        logprob, top_logprobs = extract_top_logprobs(
+            logprobs=out.logprobs,
+            tokenizer=tokenizer,
+            top_k=5,
+            selected_token=out.token,
+        )
+
         yield GenerationResponse(
-            text=out.text,
+            text=text,
             token=out.token,
-            finish_reason=cast(FinishReason | None, out.finish_reason),
+            logprob=logprob,
+            top_logprobs=top_logprobs,
+            finish_reason=finish_reason,
             stats=stats,
         )
 
-        if out.finish_reason is not None:
+        if is_done:
             break
 
+        # Limit accumulated_text to what's needed for stop sequence detection
+        if max_stop_len > 0 and len(accumulated_text) > max_stop_len:
+            accumulated_text = accumulated_text[-max_stop_len:]
+
         # TODO: Do we want an mx_barrier?

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

@@ -2,7 +2,9 @@ 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
 
@@ -40,10 +42,9 @@ import mlx.nn as nn
 from mlx_lm.utils import load_model
 from pydantic import RootModel
 
-from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.worker.instances import (
     BoundInstance,
     MlxJacclInstance,
@@ -57,8 +58,6 @@ from exo.shared.types.worker.shards import (
 from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
-    TimeoutCallback,
-    eval_with_timeout,
     pipeline_auto_parallel,
     tensor_auto_parallel,
 )
@@ -75,7 +74,7 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     return Memory.from_float_kb(
         (model_shard_meta.end_layer - model_shard_meta.start_layer)
         / model_shard_meta.n_layers
-        * model_shard_meta.model_card.storage_size.in_kb
+        * model_shard_meta.model_meta.storage_size.in_kb
         / (
             1
             if isinstance(model_shard_meta, PipelineShardMetadata)
@@ -88,6 +87,41 @@ 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(
@@ -151,28 +185,22 @@ def mlx_distributed_init(
                 group = mx.distributed.init(backend="ring", strict=True)
 
             case MlxJacclInstance(
-                jaccl_devices=jaccl_devices, jaccl_coordinators=jaccl_coordinators
+                ibv_devices=ibv_devices, jaccl_coordinators=jaccl_coordinators
             ):
-                assert all(
-                    jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
-                )
                 # Use RDMA connectivity matrix
                 coordination_file = (
                     f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
                 )
-                jaccl_devices_json = json.dumps(jaccl_devices)
+                ibv_devices_json = json.dumps(ibv_devices)
 
                 with open(coordination_file, "w") as f:
-                    _ = f.write(jaccl_devices_json)
+                    _ = f.write(ibv_devices_json)
 
                 jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
 
-                # TODO: update once upstream fixes
-                logger.info(
-                    f"rank {rank} MLX_JACCL_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
-                )
+                logger.info(f"rank {rank} MLX_IBV_DEVICES: {ibv_devices_json}")
                 logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
-                os.environ["MLX_JACCL_DEVICES"] = coordination_file
+                os.environ["MLX_IBV_DEVICES"] = coordination_file
                 os.environ["MLX_RANK"] = str(rank)
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
@@ -206,7 +234,7 @@ def load_mlx_items(
 ) -> tuple[Model, TokenizerWrapper]:
     if group is None:
         logger.info(f"Single device used for {bound_instance.instance}")
-        model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
+        model_path = build_model_path(bound_instance.bound_shard.model_meta.model_id)
         start_time = time.perf_counter()
         model, _ = load_model(model_path, strict=True)
         end_time = time.perf_counter()
@@ -234,7 +262,7 @@ def shard_and_load(
     group: Group,
     on_timeout: TimeoutCallback | None = None,
 ) -> tuple[nn.Module, TokenizerWrapper]:
-    model_path = build_model_path(shard_metadata.model_card.model_id)
+    model_path = build_model_path(shard_metadata.model_meta.model_id)
 
     model, _ = load_model(model_path, lazy=True, strict=False)
     logger.debug(model)
@@ -261,6 +289,14 @@ def shard_and_load(
 
     logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")
 
+    match shard_metadata:
+        case TensorShardMetadata():
+            logger.info(f"loading model from {model_path} with tensor parallelism")
+            model = tensor_auto_parallel(model, group)
+        case PipelineShardMetadata():
+            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)
@@ -269,15 +305,7 @@ def shard_and_load(
         f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
         f"(model size: {model_size_gb:.1f}GB)"
     )
-
-    match shard_metadata:
-        case TensorShardMetadata():
-            logger.info(f"loading model from {model_path} with tensor parallelism")
-            model = tensor_auto_parallel(model, group, timeout_seconds, on_timeout)
-        case PipelineShardMetadata():
-            logger.info(f"loading model from {model_path} with pipeline parallelism")
-            model = pipeline_auto_parallel(model, group, shard_metadata)
-            eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
+    eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
 
     # TODO: Do we need this?
     mx.eval(model)
@@ -293,7 +321,7 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
-    return load_tokenizer_for_model_id(shard_metadata.model_card.model_id, model_path)
+    return load_tokenizer_for_model_id(shard_metadata.model_meta.model_id, model_path)
 
 
 def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
@@ -312,9 +340,6 @@ def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
     model_id_lower = model_id.lower()
     if "kimi-k2" in model_id_lower:
         return [163586]
-    elif "glm-4.7-flash" in model_id_lower:
-        # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|>
-        return [154820, 154827, 154829]
     elif "glm" in model_id_lower:
         return [151336, 151329, 151338]
     return None
@@ -364,51 +389,59 @@ def load_tokenizer_for_model_id(model_id: str, model_path: Path) -> TokenizerWra
 
 def apply_chat_template(
     tokenizer: TokenizerWrapper,
-    chat_task_data: ChatCompletionTaskParams,
+    task_params: ResponsesRequest,
 ) -> str:
-    # Now we can properly access the messages
-    messages = chat_task_data.messages
+    """Convert ResponsesRequest to a chat template prompt.
 
+    Converts the internal format (input + instructions) to a messages list
+    that can be processed by the tokenizer's chat template.
+    """
     formatted_messages: list[dict[str, Any]] = []
-    for message in messages:
-        if isinstance(message.content, ChatCompletionMessageText):
-            message.content = message.content.text
-        if isinstance(message.content, list):
-            if len(message.content) == 0:
-                logger.warning("Received prompt with no content, skipping")
-                continue
 
-            message.content = "\n".join(c.text for c in message.content).strip()
-        if message.content is None and message.thinking is None:
-            continue
-
-        # Null values are not valid when applying templates in tokenizer
+    # Add system message (instructions) if present
+    if task_params.instructions:
         formatted_messages.append(
-            {k: v for k, v in message.model_dump().items() if v is not None}  # type: ignore
+            {"role": "system", "content": task_params.instructions}
         )
 
-    prompt: str = tokenizer.apply_chat_template(
-        formatted_messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        tools=chat_task_data.tools,
-    )
+    # Convert input to messages
+    if isinstance(task_params.input, str):
+        # Simple string input becomes a single user message
+        formatted_messages.append({"role": "user", "content": task_params.input})
+    else:
+        # List of InputMessage
+        for msg in task_params.input:
+            if not msg.content:
+                logger.warning("Received message with empty content, skipping")
+                continue
+            formatted_messages.append({"role": msg.role, "content": msg.content})
+
+    # Use continue_final_message when continuing from prefix (e.g., regenerate from token)
+    # This keeps the final assistant message open without EOS tokens
+    # Note: explicitly set add_generation_prompt=False when using continue_final_message
+    # because some tokenizers (e.g., Kimi) default add_generation_prompt=True
+    prompt: str
+    if task_params.continue_from_prefix:
+        prompt = tokenizer.apply_chat_template(
+            formatted_messages,
+            tokenize=False,
+            continue_final_message=True,
+            add_generation_prompt=False,
+            tools=task_params.tools,
+        )
+    else:
+        prompt = tokenizer.apply_chat_template(
+            formatted_messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            tools=task_params.tools,
+        )
 
     logger.info(prompt)
 
     return prompt
 
 
-def detect_thinking_prompt_suffix(prompt: str, tokenizer: TokenizerWrapper) -> bool:
-    """
-    Detect if prompt ends with a thinking opening tag that should be
-    prepended to the output stream.
-    """
-    think_token = tokenizer.think_start
-
-    return think_token is not None and prompt.rstrip().endswith(think_token)
-
-
 class NullKVCache(KVCache):
     """
     A KVCache that pretends to exist but holds zero tokens.

src/exo/worker/main.py

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

src/exo/worker/plan.py

@@ -2,8 +2,8 @@
 
 from collections.abc import Mapping, Sequence
 
-from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import NodeId
+from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -114,7 +114,7 @@ def _model_needs_download(
     download_status: Mapping[ModelId, DownloadProgress],
 ) -> DownloadModel | None:
     for runner in runners.values():
-        model_id = runner.bound_instance.bound_shard.model_card.model_id
+        model_id = runner.bound_instance.bound_shard.model_meta.model_id
         if isinstance(runner.status, RunnerIdle) and (
             model_id not in download_status
             or not isinstance(
@@ -191,7 +191,7 @@ def _load_model(
             nid in global_download_status
             and any(
                 isinstance(dp, DownloadCompleted)
-                and dp.shard_metadata.model_card.model_id == shard_assignments.model_id
+                and dp.shard_metadata.model_meta.model_id == shard_assignments.model_id
                 for dp in global_download_status[nid]
             )
             for nid in shard_assignments.node_to_runner

src/exo/worker/runner/bootstrap.py

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

src/exo/worker/runner/runner.py

@@ -4,7 +4,6 @@ from functools import cache
 
 import mlx.core as mx
 from mlx_lm.models.gpt_oss import Model as GptOssModel
-from mlx_lm.tokenizer_utils import TokenizerWrapper
 from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     HarmonyEncodingName,
     Role,
@@ -12,7 +11,6 @@ from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     load_harmony_encoding,
 )
 
-from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.events import (
     ChunkGenerated,
@@ -21,6 +19,7 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -51,8 +50,6 @@ from exo.shared.types.worker.runners import (
 from exo.utils.channels import MpReceiver, MpSender
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
-    apply_chat_template,
-    detect_thinking_prompt_suffix,
     initialize_mlx,
     load_mlx_items,
     mlx_force_oom,
@@ -175,33 +172,21 @@ def main(
                     )
                     assert model
                     assert tokenizer
-                    assert task_params.messages[0].content is not None
 
                     try:
-                        _check_for_debug_prompts(task_params.messages[0].content)
-
-                        # Build prompt once - used for both generation and thinking detection
-                        prompt = apply_chat_template(tokenizer, task_params)
+                        _check_for_debug_prompts(task_params)
 
                         # Generate responses using the actual MLX generation
                         mlx_generator = mlx_generate(
                             model=model,
                             tokenizer=tokenizer,
                             task=task_params,
-                            prompt=prompt,
                         )
 
                         # GPT-OSS specific parsing to match other model formats.
                         if isinstance(model, GptOssModel):
                             mlx_generator = parse_gpt_oss(mlx_generator)
 
-                        # For other thinking models (GLM, etc.), check if we need to
-                        # prepend the thinking tag that was consumed by the chat template
-                        if detect_thinking_prompt_suffix(prompt, tokenizer):
-                            mlx_generator = parse_thinking_models(
-                                mlx_generator, tokenizer
-                            )
-
                         # TODO: Add tool call parser here
 
                         for response in mlx_generator:
@@ -213,9 +198,11 @@ def main(
                                                 command_id=command_id,
                                                 chunk=TokenChunk(
                                                     idx=response.token,
-                                                    model=shard_metadata.model_card.model_id,
+                                                    model=shard_metadata.model_meta.model_id,
                                                     text=response.text,
                                                     token_id=response.token,
+                                                    logprob=response.logprob,
+                                                    top_logprobs=response.top_logprobs,
                                                     finish_reason=response.finish_reason,
                                                     stats=response.stats,
                                                 ),
@@ -230,7 +217,7 @@ def main(
                                     command_id=command_id,
                                     chunk=TokenChunk(
                                         idx=0,
-                                        model=shard_metadata.model_card.model_id,
+                                        model=shard_metadata.model_meta.model_id,
                                         text="",
                                         token_id=0,
                                         finish_reason="error",
@@ -307,44 +294,28 @@ def parse_gpt_oss(
             break
 
 
-def parse_thinking_models(
-    responses: Generator[GenerationResponse],
-    tokenizer: TokenizerWrapper,
-) -> Generator[GenerationResponse]:
-    """
-    For models that inject thinking tags in the prompt (like GLM-4.7),
-    prepend the thinking tag to the output stream so the frontend
-    can properly parse thinking content.
-    """
-    first = True
-    for response in responses:
-        if first:
-            first = False
-            yield response.model_copy(
-                update={
-                    "text": tokenizer.think_start,
-                    "token": tokenizer.think_start_id,  # type: ignore
-                }
-            )
-        yield response
-
-
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
 EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
 EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
 
 
-def _check_for_debug_prompts(
-    prompt: str | ChatCompletionMessageText | list[ChatCompletionMessageText],
-):
-    if isinstance(prompt, list):
-        if len(prompt) == 0:
-            logger.debug("Empty message prompt received in debug prompt")
-            return
-        prompt = prompt[0]
+def _check_for_debug_prompts(task_params: ResponsesRequest) -> None:
+    """Check for debug prompt triggers in the input.
 
-    if isinstance(prompt, ChatCompletionMessageText):
-        prompt = prompt.text
+    Extracts the first user input text and checks for debug triggers.
+    """
+    prompt: str
+    if isinstance(task_params.input, str):
+        prompt = task_params.input
+    else:
+        # List of InputMessage - get first message content
+        if len(task_params.input) == 0:
+            logger.debug("Empty message list in debug prompt check")
+            return
+        prompt = task_params.input[0].content
+
+    if not prompt:
+        return
 
     if EXO_RUNNER_MUST_FAIL in prompt:
         logger.info("raising exception")

src/exo/worker/tests/constants.py

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

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

@@ -1,8 +1,8 @@
 from dataclasses import dataclass, field
 
-from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.tasks import BaseTask, TaskId
 from exo.shared.types.worker.instances import (
     BoundInstance,
@@ -32,8 +32,9 @@ def get_pipeline_shard_metadata(
     model_id: ModelId, device_rank: int, world_size: int = 1
 ) -> ShardMetadata:
     return PipelineShardMetadata(
-        model_card=ModelCard(
+        model_meta=ModelMetadata(
             model_id=model_id,
+            pretty_name=str(model_id),
             storage_size=Memory.from_mb(100000),
             n_layers=32,
             hidden_size=2048,

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

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

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

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

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

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

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

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

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

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

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

@@ -1,7 +1,7 @@
 from typing import cast
 
 import exo.worker.plan as plan_mod
-from exo.shared.types.api import ChatCompletionTaskParams
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.tasks import ChatCompletion, Task, TaskId, TaskStatus
 from exo.shared.types.worker.instances import BoundInstance, InstanceId
 from exo.shared.types.worker.runners import (
@@ -59,7 +59,7 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
         instance_id=INSTANCE_1_ID,
         task_status=TaskStatus.Pending,
         command_id=COMMAND_1_ID,
-        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+        task_params=ResponsesRequest(model=MODEL_A_ID, input=""),
     )
 
     result = plan_mod.plan(
@@ -107,7 +107,7 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
         instance_id=INSTANCE_1_ID,
         task_status=TaskStatus.Pending,
         command_id=COMMAND_1_ID,
-        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+        task_params=ResponsesRequest(model=MODEL_A_ID, input=""),
     )
 
     result = plan_mod.plan(
@@ -152,7 +152,7 @@ def test_plan_does_not_forward_tasks_for_other_instances():
         instance_id=other_instance_id,
         task_status=TaskStatus.Pending,
         command_id=COMMAND_1_ID,
-        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+        task_params=ResponsesRequest(model=MODEL_A_ID, input=""),
     )
 
     result = plan_mod.plan(
@@ -201,7 +201,7 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
         instance_id=INSTANCE_1_ID,
         task_status=TaskStatus.Complete,
         command_id=COMMAND_1_ID,
-        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+        task_params=ResponsesRequest(model=MODEL_A_ID, input=""),
     )
 
     other_task_id = TaskId("other-task")

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

@@ -5,7 +5,6 @@ from typing import Callable
 import pytest
 
 import exo.worker.runner.runner as mlx_runner
-from exo.shared.types.api import ChatCompletionMessage
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.events import (
     ChunkGenerated,
@@ -14,9 +13,9 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.tasks import (
     ChatCompletion,
-    ChatCompletionTaskParams,
     ConnectToGroup,
     LoadModel,
     Shutdown,
@@ -85,11 +84,11 @@ SHUTDOWN_TASK = Shutdown(
     runner_id=RUNNER_1_ID,
 )
 
-CHAT_PARAMS = ChatCompletionTaskParams(
+CHAT_PARAMS = ResponsesRequest(
     model=str(MODEL_A_ID),
-    messages=[ChatCompletionMessage(role="user", content="hello")],
+    input="hello",
     stream=True,
-    max_tokens=4,
+    max_output_tokens=4,
     temperature=0.0,
 )
 
@@ -114,10 +113,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, 1)))
     monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
-    # Mock apply_chat_template since we're using a fake tokenizer (integer 1).
-    # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None.
-    monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt"))
-    monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
 
     def fake_generate(*_1: object, **_2: object):
         yield GenerationResponse(token=0, text="hi", finish_reason="stop")

src/exo/worker/utils/__init__.py

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

src/exo/worker/utils/macmon.py

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

src/exo/worker/utils/net_profile.py

@@ -1,5 +1,3 @@
-from collections.abc import Mapping
-
 import anyio
 import httpx
 from anyio import create_task_group
@@ -7,7 +5,6 @@ from loguru import logger
 
 from exo.shared.topology import Topology
 from exo.shared.types.common import NodeId
-from exo.shared.types.profiling import NodeNetworkInfo
 
 REACHABILITY_ATTEMPTS = 3
 
@@ -21,9 +18,8 @@ async def check_reachability(
     """Check if a node is reachable at the given IP and verify its identity."""
     if ":" in target_ip:
         # TODO: use real IpAddress types
-        url = f"http://[{target_ip}]:52415/node_id"
-    else:
-        url = f"http://{target_ip}:52415/node_id"
+        target_ip = f"[{target_ip}]"
+    url = f"http://{target_ip}:52415/node_id"
 
     remote_node_id = None
     last_error = None
@@ -77,9 +73,7 @@ async def check_reachability(
 
 
 async def check_reachable(
-    topology: Topology,
-    self_node_id: NodeId,
-    node_network: Mapping[NodeId, NodeNetworkInfo],
+    topology: Topology, self_node_id: NodeId
 ) -> dict[NodeId, set[str]]:
     """Check which nodes are reachable and return their IPs."""
 
@@ -97,16 +91,16 @@ async def check_reachable(
         httpx.AsyncClient(timeout=timeout, limits=limits) as client,
         create_task_group() as tg,
     ):
-        for node_id in topology.list_nodes():
-            if node_id not in node_network:
+        for node in topology.list_nodes():
+            if not node.node_profile:
                 continue
-            if node_id == self_node_id:
+            if node.node_id == self_node_id:
                 continue
-            for iface in node_network[node_id].interfaces:
+            for iface in node.node_profile.network_interfaces:
                 tg.start_soon(
                     check_reachability,
                     iface.ip_address,
-                    node_id,
+                    node.node_id,
                     reachable,
                     client,
                 )

src/exo/worker/utils/profile.py

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

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

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

tests/headless_runner.py

@@ -13,10 +13,10 @@ from pydantic import BaseModel
 
 from exo.shared.logging import InterceptLogger, logger_setup
 from exo.shared.models.model_cards import MODEL_CARDS, ModelId
-from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
 from exo.shared.types.commands import CommandId
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.events import Event
+from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -34,8 +34,7 @@ from exo.shared.types.worker.instances import (
 )
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
-from exo.utils.channels import MpReceiver, MpSender, channel, mp_channel
-from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
+from exo.utils.channels import MpReceiver, MpSender, mp_channel
 from exo.worker.download.impl_shard_downloader import (
     build_full_shard,
     exo_shard_downloader,
@@ -66,7 +65,6 @@ async def main():
     app = FastAPI()
     app.post("/ring")(ring_backend)
     app.post("/jaccl")(jaccl_backend)
-    app.post("/tb_detection")(tb_detection)
     shutdown = anyio.Event()
     await serve(
         app,  # type: ignore
@@ -78,15 +76,6 @@ async def main():
     shutdown.set()
 
 
-async def tb_detection():
-    send, recv = channel[GatheredInfo]()
-    ig = InfoGatherer(send)
-    with anyio.move_on_after(1):
-        await ig._monitor_system_profiler_thunderbolt_data()  # pyright: ignore[reportPrivateUsage]
-    with recv:
-        return recv.collect()
-
-
 async def assert_downloads():
     sd = exo_shard_downloader()
     # await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-0.6b"].model_id))
@@ -135,7 +124,7 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
     else:
         raise ValueError(f"{hn} not in {test.devs}")
 
-    card = MODEL_CARDS[test.model_id]
+    meta = MODEL_CARDS[test.model_id].metadata
     instance = MlxRingInstance(
         instance_id=iid,
         ephemeral_port=52416,
@@ -145,15 +134,15 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
             node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
             runner_to_shard={
                 RunnerId(test.devs[i][0]): PipelineShardMetadata(
-                    model_card=card,
+                    model_meta=meta,
                     device_rank=i,
                     world_size=world_size,
-                    start_layer=(card.n_layers // world_size) * i,
+                    start_layer=(meta.n_layers // world_size) * i,
                     end_layer=min(
-                        card.n_layers, (card.n_layers // world_size) * (i + 1)
+                        meta.n_layers, (meta.n_layers // world_size) * (i + 1)
                     ),
-                    n_layers=min(card.n_layers, (card.n_layers // world_size) * (i + 1))
-                    - (card.n_layers // world_size) * i,
+                    n_layers=min(meta.n_layers, (meta.n_layers // world_size) * (i + 1))
+                    - (meta.n_layers // world_size) * i,
                 )
                 for i in range(world_size)
             },
@@ -180,16 +169,10 @@ async def execute_test(test: Tests, instance: Instance, hn: str):
             send.send(StartWarmup(instance_id=iid))
             send.send(
                 ChatCompletion(
-                    task_params=ChatCompletionTaskParams(
+                    task_params=ResponsesRequest(
                         model=test.model_id,
-                        messages=[
-                            ChatCompletionMessage(
-                                role="system", content="You are a helpful assistant"
-                            ),
-                            ChatCompletionMessage(
-                                role="user", content="What is the capital of France?"
-                            ),
-                        ],
+                        instructions="You are a helpful assistant",
+                        input="What is the capital of France?",
                     ),
                     command_id=CommandId("yo"),
                     instance_id=iid,
@@ -220,16 +203,16 @@ async def jaccl_backend(test: Tests):
             break
     else:
         raise ValueError(f"{weird_hn} not in {test.devs}")
-    return await execute_test(test, jaccl_instance(test, iid), hn)
+    return await execute_test(test, jaccl_instance(test, iid, hn), hn)
 
 
-def jaccl_instance(test: Tests, iid: InstanceId):
-    card = MODEL_CARDS[test.model_id]
+def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
+    meta = MODEL_CARDS[test.model_id].metadata
     world_size = len(test.devs)
 
     return MlxJacclInstance(
         instance_id=iid,
-        jaccl_devices=[[None, "rdma_en3"], ["rdma_en3", None]],
+        ibv_devices=[[None, "rdma_en3"], ["rdma_en3", None]],
         # rank 0 is always coordinator
         jaccl_coordinators={
             NodeId(host[0]): test.devs[0][1] + ":52416" for host in test.devs
@@ -239,12 +222,12 @@ def jaccl_instance(test: Tests, iid: InstanceId):
             node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
             runner_to_shard={
                 RunnerId(test.devs[i][0]): TensorShardMetadata(
-                    model_card=card,
+                    model_meta=meta,
                     device_rank=i,
                     world_size=world_size,
-                    start_layer=card.n_layers,
-                    end_layer=card.n_layers,
-                    n_layers=card.n_layers,
+                    start_layer=meta.n_layers,
+                    end_layer=meta.n_layers,
+                    n_layers=meta.n_layers,
                 )
                 for i in range(world_size)
             },

tmp/kill_bridge_plist.sh

@@ -1,84 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-PREFS="${PREFS:-/Library/Preferences/SystemConfiguration/preferences.plist}"
-
-tmpdir="$(mktemp -d)"
-trap 'rm -rf "$tmpdir"' EXIT
-injson="$tmpdir/in.json"
-outjson="$tmpdir/out.json"
-plutil -convert json -o "$injson" "$PREFS"
-
-perl -Mstrict -Mwarnings -MJSON::PP -e '
-  my ($in, $out) = @ARGV;
-
-  open my $fh, "<", $in or die "open $in: $!";
-  local $/;
-  my $txt = <$fh>;
-  close $fh;
-
-  my $json = JSON::PP->new->utf8->relaxed(1);
-  my $d = $json->decode($txt);
-
-  if (ref($d->{VirtualNetworkInterfaces}) eq "HASH"
-      && ref($d->{VirtualNetworkInterfaces}{Bridge}) eq "HASH") {
-    delete $d->{VirtualNetworkInterfaces}{Bridge}{bridge0};
-  }
-
-  my @bridge_svcs;
-  if (ref($d->{NetworkServices}) eq "HASH") {
-    for my $k (keys %{ $d->{NetworkServices} }) {
-      my $svc = $d->{NetworkServices}{$k};
-      next unless ref($svc) eq "HASH";
-      my $iface = $svc->{Interface};
-      next unless ref($iface) eq "HASH";
-      my $dev = $iface->{DeviceName};
-      if (defined $dev && $dev eq "bridge0") {
-        push @bridge_svcs, $k;
-      }
-    }
-    delete @{ $d->{NetworkServices} }{ @bridge_svcs } if @bridge_svcs;
-  }
-
-  my %is_bridge = map { $_ => 1 } @bridge_svcs;
-
-  if (ref($d->{Sets}) eq "HASH") {
-    for my $setk (keys %{ $d->{Sets} }) {
-      my $set = $d->{Sets}{$setk};
-      next unless ref($set) eq "HASH";
-      my $net = $set->{Network};
-      next unless ref($net) eq "HASH";
-
-      if (ref($net->{Interface}) eq "HASH") {
-        delete $net->{Interface}{bridge0};
-      }
-
-      if (ref($net->{Service}) eq "HASH" && @bridge_svcs) {
-        for my $svc (@bridge_svcs) {
-          delete $net->{Service}{$svc};
-        }
-      }
-
-      my $g = $net->{Global};
-      if (ref($g) eq "HASH"
-          && ref($g->{IPv4}) eq "HASH"
-          && ref($g->{IPv4}{ServiceOrder}) eq "ARRAY"
-          && @bridge_svcs) {
-
-        my @so = @{ $g->{IPv4}{ServiceOrder} };
-        @so = grep { !defined($_) || !$is_bridge{$_} } @so;
-        $g->{IPv4}{ServiceOrder} = \@so;
-      }
-    }
-  }
-
-  open my $oh, ">", $out or die "open $out: $!";
-  print $oh JSON::PP->new->utf8->canonical(1)->pretty(1)->encode($d);
-  close $oh;
-' "$injson" "$outjson"
-
-# Convert JSON -> plist (write back as binary1; change to xml1 if you prefer)
-plutil -convert xml1 -o "$PREFS" "$outjson"
-
-# Ask configd to reload SystemConfiguration state
-killall -HUP configd 2>/dev/null || true