formatting

code quality
Merge main into sami/flash
2026-01-20 03:51:14 -05:00 · 2026-01-20 07:54:52 +05:00 · 2026-01-20 07:47:14 +05:00 · 2026-01-20 07:31:22 +05:00 · 2026-01-20 07:21:41 +05:00 · 2026-01-20 06:53:43 +05:00
97 changed files with 5535 additions and 4334 deletions
--- a/README.md
+++ b/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,6 +163,24 @@ This starts the exo dashboard and API at http://localhost:52415/

 **Important note for Linux users:** Currently, exo runs on CPU on Linux. GPU support for Linux platforms is under development. If you'd like to see support for your specific Linux hardware, please [search for existing feature requests](https://github.com/exo-explore/exo/issues) or create a new one.

+**Configuration Options:**
+
+- `--no-worker`: Run exo without the worker component. Useful for coordinator-only nodes that handle networking and orchestration but don't execute inference tasks. This is helpful for machines without sufficient GPU resources but with good network connectivity.
+
+  ```bash
+  uv run exo --no-worker
+  ```
+
+**File Locations (Linux):**
+
+exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux:
+
+- **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
+- **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
+- **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
+
+You can override these locations by setting the corresponding XDG environment variables.
+
 ### macOS App

 exo ships a macOS app that runs in the background on your Mac.
@@ -175,6 +193,19 @@ Download the latest build here: [EXO-latest.dmg](https://assets.exolabs.net/EXO-

 The app will ask for permission to modify system settings and install a new Network profile. Improvements to this are being worked on.

+**Custom Namespace for Cluster Isolation:**
+
+The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. This is configured through the `EXO_LIBP2P_NAMESPACE` setting:
+
+- **Use cases**:
+  - Running multiple separate exo clusters on the same network
+  - Isolating development/testing clusters from production clusters
+  - Preventing accidental cluster joining
+
+- **Configuration**: Access this setting in the app's Advanced settings (or set the `EXO_LIBP2P_NAMESPACE` environment variable when running from source)
+
+The namespace is logged on startup for debugging purposes.
+
 #### Uninstalling the macOS App

 The recommended way to uninstall is through the app itself: click the menu bar icon → Advanced → Uninstall. This cleanly removes all system components.
@@ -321,6 +352,52 @@ For further details, see:

 ---

+## Benchmarking
+
+The `exo-bench` tool measures model prefill and token generation speed across different placement configurations. This helps you optimize model performance and validate improvements.
+
+**Prerequisites:**
+- Nodes should be running with `uv run exo` before benchmarking
+- The tool uses the `/bench/chat/completions` endpoint
+
+**Basic usage:**
+
+```bash
+uv run bench/exo_bench.py \
+  --model llama-3.2-1b \
+  --pp 128,256,512 \
+  --tg 128,256
+```
+
+**Key parameters:**
+
+- `--model`: Model to benchmark (short ID or HuggingFace ID)
+- `--pp`: Prompt size hints (comma-separated integers)
+- `--tg`: Generation lengths (comma-separated integers)
+- `--max-nodes`: Limit placements to N nodes (default: 4)
+- `--instance-meta`: Filter by `ring`, `jaccl`, or `both` (default: both)
+- `--sharding`: Filter by `pipeline`, `tensor`, or `both` (default: both)
+- `--repeat`: Number of repetitions per configuration (default: 1)
+- `--warmup`: Warmup runs per placement (default: 0)
+- `--json-out`: Output file for results (default: bench/results.json)
+
+**Example with filters:**
+
+```bash
+uv run bench/exo_bench.py \
+  --model llama-3.2-1b \
+  --pp 128,512 \
+  --tg 128 \
+  --max-nodes 2 \
+  --sharding tensor \
+  --repeat 3 \
+  --json-out my-results.json
+```
+
+The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
+
+---
+
 ## Hardware Accelerator Support

 On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working on extending hardware accelerator support. If you'd like support for a new hardware platform, please [search for an existing feature request](https://github.com/exo-explore/exo/issues) and add a thumbs up so we know what hardware is important to the community.
@@ -329,4 +406,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working

 ## Contributing

-See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
--- a/TODO.md
+++ b/TODO.md
@@ -19,6 +19,7 @@
 25. Rethink retry logic
 26. Task cancellation. When API http request gets cancelled, it should cancel corresponding task.
 27. Log cleanup - per-module log filters and default to DEBUG log levels
+28. Validate RDMA connections with ibv_devinfo in the info gatherer

 Potential refactors:

--- a/app/EXO/EXO/Services/NetworkSetupHelper.swift
+++ b/app/EXO/EXO/Services/NetworkSetupHelper.swift
@@ -6,7 +6,7 @@ enum NetworkSetupHelper {
    private static let logger = Logger(subsystem: "io.exo.EXO", category: "NetworkSetup")
    private static let daemonLabel = "io.exo.networksetup"
    private static let scriptDestination =
-        "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
+        "/Library/Application Support/EXO/disable_bridge.sh"
    private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist"
    private static let requiredStartInterval: Int = 1791

@@ -28,35 +28,6 @@ enum NetworkSetupHelper {
        # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist
        /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true

-        networksetup -listlocations | grep -q exo || {
-          networksetup -createlocation exo
-        }
-
-        networksetup -switchtolocation exo
-        networksetup -listallhardwareports \\
-          | awk -F': ' '/Hardware Port: / {print $2}' \\
-          | while IFS=":" read -r name; do
-              case "$name" in
-                "Ethernet Adapter"*)
-                        ;;
-                "Thunderbolt Bridge")
-                        ;;
-                "Thunderbolt "*)
-                  networksetup -listallnetworkservices \\
-                    | grep -q "EXO $name" \\
-                      || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \\
-                      || continue
-                  networksetup -setdhcp "EXO $name"
-                        ;;
-                *)
-                  networksetup -listallnetworkservices \\
-                    | grep -q "$name" \\
-                      || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \\
-                      || continue
-                        ;;
-              esac
-            done
-
        networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
          networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
        } || true
@@ -141,6 +112,13 @@ enum NetworkSetupHelper {
        let scriptExists = manager.fileExists(atPath: scriptDestination)
        let plistExists = manager.fileExists(atPath: plistDestination)
        guard scriptExists, plistExists else { return false }
+        guard
+            let installedScript = try? String(contentsOfFile: scriptDestination, encoding: .utf8),
+            installedScript.trimmingCharacters(in: .whitespacesAndNewlines)
+                == setupScript.trimmingCharacters(in: .whitespacesAndNewlines)
+        else {
+            return false
+        }
        guard
            let data = try? Data(contentsOf: URL(fileURLWithPath: plistDestination)),
            let plist = try? PropertyListSerialization.propertyList(
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -16,9 +16,6 @@ from urllib.parse import urlencode
 from loguru import logger
 from transformers import AutoTokenizer

-from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.types.memory import Memory
-

 class ExoHttpError(RuntimeError):
    def __init__(self, status: int, reason: str, body_preview: str):
@@ -490,17 +487,17 @@ def main() -> int:
                logger.debug(f"  warmup {i + 1}/{args.warmup} done")

            for pp in pp_list:
-                if (
-                    pp * n_nodes > 2048
-                    and "ring" in instance_meta.lower()
-                    and "tensor" in sharding.lower()
-                ):
-                    model_card = MODEL_CARDS[short_id]
-                    if model_card.metadata.storage_size > Memory.from_gb(10):
-                        logger.info(
-                            f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
-                        )
-                        continue
+                # if (
+                #     pp * n_nodes > 2048
+                #     and "ring" in instance_meta.lower()
+                #     and "tensor" in sharding.lower()
+                # ):
+                #     model_card = MODEL_CARDS[short_id]
+                #     if model_card.metadata.storage_size > Memory.from_gb(10):
+                #         logger.info(
+                #             f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
+                #         )
+                #         continue
                for tg in tg_list:
                    runs: list[dict[str, Any]] = []
                    for r in range(args.repeat):
--- a/dashboard/src/lib/components/ChatMessages.svelte
+++ b/dashboard/src/lib/components/ChatMessages.svelte
@@ -1,16 +1,14 @@
 <script lang="ts">
-	import {
-		messages,
-		currentResponse,
+	import { 
+		messages, 
+		currentResponse, 
 		isLoading,
 		deleteMessage,
 		editAndRegenerate,
-		regenerateLastResponse,
-		regenerateFromToken
+		regenerateLastResponse
 	} 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;
@@ -97,23 +95,6 @@
 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,
@@ -385,17 +366,7 @@ function isThinkingExpanded(messageId: string): boolean {
 									</div>
 								{/if}
 								<div class="text-xs text-foreground">
-									{#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}
+									<MarkdownContent content={message.content || (loading ? response : '')} />
 									{#if loading && !message.content}
 										<span class="inline-block w-2 h-4 bg-exo-yellow/70 ml-1 cursor-blink"></span>
 									{/if}
@@ -448,20 +419,7 @@ 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)}
--- a/dashboard/src/lib/components/MarkdownContent.svelte
+++ b/dashboard/src/lib/components/MarkdownContent.svelte
@@ -53,62 +53,285 @@
 	marked.use({ renderer });

 	/**
-	 * Preprocess LaTeX: convert \(...\) to $...$ and \[...\] to $$...$$
-	 * Also protect code blocks from LaTeX processing
+	 * Unescape HTML entities that marked may have escaped
+	 */
+	function unescapeHtmlEntities(text: string): string {
+		return text
+			.replace(/&lt;/g, '<')
+			.replace(/&gt;/g, '>')
+			.replace(/&amp;/g, '&')
+			.replace(/&quot;/g, '"')
+			.replace(/&#39;/g, "'");
+	}
+
+	// Storage for math expressions extracted before markdown processing
+	const mathExpressions: Map<string, { content: string; displayMode: boolean }> = new Map();
+	let mathCounter = 0;
+
+	// Storage for HTML snippets that need protection from markdown
+	const htmlSnippets: Map<string, string> = new Map();
+	let htmlCounter = 0;
+
+	// Use alphanumeric placeholders that won't be interpreted as HTML tags
+	const MATH_PLACEHOLDER_PREFIX = 'MATHPLACEHOLDER';
+	const CODE_PLACEHOLDER_PREFIX = 'CODEPLACEHOLDER';
+	const HTML_PLACEHOLDER_PREFIX = 'HTMLPLACEHOLDER';
+
+	/**
+	 * Preprocess LaTeX: extract math, handle LaTeX document commands, and protect content
 	 */
 	function preprocessLaTeX(text: string): string {
-		// Protect code blocks
+		// Reset storage
+		mathExpressions.clear();
+		mathCounter = 0;
+		htmlSnippets.clear();
+		htmlCounter = 0;
+
+		// Protect code blocks first
 		const codeBlocks: string[] = [];
 		let processed = text.replace(/```[\s\S]*?```|`[^`]+`/g, (match) => {
 			codeBlocks.push(match);
-			return `<<CODE_${codeBlocks.length - 1}>>`;
+			return `${CODE_PLACEHOLDER_PREFIX}${codeBlocks.length - 1}END`;
 		});

-		// Convert \(...\) to $...$
-		processed = processed.replace(/\\\((.+?)\\\)/g, '$$$1$');
-		
-		// Convert \[...\] to $$...$$
-		processed = processed.replace(/\\\[([\s\S]*?)\\\]/g, '$$$$$1$$$$');
+		// Remove LaTeX document commands
+		processed = processed.replace(/\\documentclass(\[[^\]]*\])?\{[^}]*\}/g, '');
+		processed = processed.replace(/\\usepackage(\[[^\]]*\])?\{[^}]*\}/g, '');
+		processed = processed.replace(/\\begin\{document\}/g, '');
+		processed = processed.replace(/\\end\{document\}/g, '');
+		processed = processed.replace(/\\maketitle/g, '');
+		processed = processed.replace(/\\title\{[^}]*\}/g, '');
+		processed = processed.replace(/\\author\{[^}]*\}/g, '');
+		processed = processed.replace(/\\date\{[^}]*\}/g, '');
+
+		// Remove \require{...} commands (MathJax-specific, not supported by KaTeX)
+		processed = processed.replace(/\$\\require\{[^}]*\}\$/g, '');
+		processed = processed.replace(/\\require\{[^}]*\}/g, '');
+
+		// Remove unsupported LaTeX commands/environments (tikzpicture, figure, center, etc.)
+		processed = processed.replace(/\\begin\{tikzpicture\}[\s\S]*?\\end\{tikzpicture\}/g, () => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, '<div class="latex-diagram-placeholder"><span class="latex-diagram-icon">📐</span><span class="latex-diagram-text">Diagram</span></div>');
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\begin\{figure\}[\s\S]*?\\end\{figure\}/g, () => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, '<div class="latex-diagram-placeholder"><span class="latex-diagram-icon">🖼️</span><span class="latex-diagram-text">Figure</span></div>');
+			htmlCounter++;
+			return placeholder;
+		});
+		// Strip center environment (layout only, no content change)
+		processed = processed.replace(/\\begin\{center\}/g, '');
+		processed = processed.replace(/\\end\{center\}/g, '');
+		// Strip other layout environments
+		processed = processed.replace(/\\begin\{flushleft\}/g, '');
+		processed = processed.replace(/\\end\{flushleft\}/g, '');
+		processed = processed.replace(/\\begin\{flushright\}/g, '');
+		processed = processed.replace(/\\end\{flushright\}/g, '');
+		processed = processed.replace(/\\label\{[^}]*\}/g, '');
+		processed = processed.replace(/\\caption\{[^}]*\}/g, '');
+
+		// Protect escaped dollar signs (e.g., \$50 should become $50, not LaTeX)
+		processed = processed.replace(/\\\$/g, 'ESCAPEDDOLLARPLACEHOLDER');
+
+		// Convert LaTeX math environments to display math (both bare and wrapped in $...$)
+		const mathEnvs = ['align', 'align\\*', 'equation', 'equation\\*', 'gather', 'gather\\*', 'multline', 'multline\\*', 'eqnarray', 'eqnarray\\*', 'array', 'matrix', 'pmatrix', 'bmatrix', 'vmatrix', 'cases'];
+		for (const env of mathEnvs) {
+			// Handle $\begin{env}...\end{env}$ (with dollar signs, possibly multiline)
+			const wrappedRegex = new RegExp(`\\$\\\\begin\\{${env}\\}(\\{[^}]*\\})?([\\s\\S]*?)\\\\end\\{${env}\\}\\$`, 'g');
+			processed = processed.replace(wrappedRegex, (_, args, content) => {
+				const cleanEnv = env.replace('\\*', '*');
+				const mathContent = `\\begin{${cleanEnv}}${args || ''}${content}\\end{${cleanEnv}}`;
+				const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+				mathExpressions.set(placeholder, { content: mathContent, displayMode: true });
+				mathCounter++;
+				return placeholder;
+			});
+
+			// Handle bare \begin{env}...\end{env} (without dollar signs)
+			const bareRegex = new RegExp(`\\\\begin\\{${env}\\}(\\{[^}]*\\})?([\\s\\S]*?)\\\\end\\{${env}\\}`, 'g');
+			processed = processed.replace(bareRegex, (_, args, content) => {
+				const cleanEnv = env.replace('\\*', '*');
+				const mathContent = `\\begin{${cleanEnv}}${args || ''}${content}\\end{${cleanEnv}}`;
+				const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+				mathExpressions.set(placeholder, { content: mathContent, displayMode: true });
+				mathCounter++;
+				return placeholder;
+			});
+		}
+
+		// Convert LaTeX proof environments to styled blocks (use placeholders for HTML)
+		processed = processed.replace(
+			/\\begin\{proof\}([\s\S]*?)\\end\{proof\}/g,
+			(_, content) => {
+				const html = `<div class="latex-proof"><div class="latex-proof-header">Proof</div><div class="latex-proof-content">${content}</div></div>`;
+				const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+				htmlSnippets.set(placeholder, html);
+				htmlCounter++;
+				return placeholder;
+			}
+		);
+
+		// Convert LaTeX theorem-like environments
+		const theoremEnvs = ['theorem', 'lemma', 'corollary', 'proposition', 'definition', 'remark', 'example'];
+		for (const env of theoremEnvs) {
+			const envRegex = new RegExp(`\\\\begin\\{${env}\\}([\\s\\S]*?)\\\\end\\{${env}\\}`, 'gi');
+			const envName = env.charAt(0).toUpperCase() + env.slice(1);
+			processed = processed.replace(envRegex, (_, content) => {
+				const html = `<div class="latex-theorem"><div class="latex-theorem-header">${envName}</div><div class="latex-theorem-content">${content}</div></div>`;
+				const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+				htmlSnippets.set(placeholder, html);
+				htmlCounter++;
+				return placeholder;
+			});
+		}
+
+		// Convert LaTeX text formatting commands (use placeholders to protect from markdown)
+		processed = processed.replace(/\\emph\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<em>${content}</em>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\textit\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<em>${content}</em>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\textbf\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<strong>${content}</strong>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\texttt\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<code class="inline-code">${content}</code>`);
+			htmlCounter++;
+			return placeholder;
+		});
+		processed = processed.replace(/\\underline\{([^}]*)\}/g, (_, content) => {
+			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
+			htmlSnippets.set(placeholder, `<u>${content}</u>`);
+			htmlCounter++;
+			return placeholder;
+		});
+
+		// Handle LaTeX line breaks and spacing
+		processed = processed.replace(/\\\\(?:\s*\n)?/g, '\n'); // \\ -> newline
+		processed = processed.replace(/\\newline/g, '\n');
+		processed = processed.replace(/\\par\b/g, '\n\n');
+		processed = processed.replace(/\\quad/g, ' ');
+		processed = processed.replace(/\\qquad/g, '  ');
+		processed = processed.replace(/~~/g, ' '); // non-breaking space
+
+		// Remove other common LaTeX commands that don't render
+		processed = processed.replace(/\\centering/g, '');
+		processed = processed.replace(/\\noindent/g, '');
+		processed = processed.replace(/\\hfill/g, '');
+		processed = processed.replace(/\\vspace\{[^}]*\}/g, '');
+		processed = processed.replace(/\\hspace\{[^}]*\}/g, ' ');
+
+		// Convert \(...\) to placeholder (display: false)
+		processed = processed.replace(/\\\(([\s\S]+?)\\\)/g, (_, content) => {
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}INLINE${mathCounter}END`;
+			mathExpressions.set(placeholder, { content, displayMode: false });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Convert \[...\] to placeholder (display: true)
+		processed = processed.replace(/\\\[([\s\S]*?)\\\]/g, (_, content) => {
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+			mathExpressions.set(placeholder, { content, displayMode: true });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Extract display math ($$...$$) BEFORE markdown processing
+		processed = processed.replace(/\$\$([\s\S]*?)\$\$/g, (_, content) => {
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
+			mathExpressions.set(placeholder, { content: content.trim(), displayMode: true });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Extract inline math ($...$) BEFORE markdown processing
+		// Allow single-line only, skip currency patterns like $5 or $50
+		processed = processed.replace(/\$([^\$\n]+?)\$/g, (match, content) => {
+			if (/^\d/.test(content.trim())) {
+				return match; // Keep as-is for currency
+			}
+			const placeholder = `${MATH_PLACEHOLDER_PREFIX}INLINE${mathCounter}END`;
+			mathExpressions.set(placeholder, { content: content.trim(), displayMode: false });
+			mathCounter++;
+			return placeholder;
+		});
+
+		// Restore escaped dollar signs
+		processed = processed.replace(/ESCAPEDDOLLARPLACEHOLDER/g, '$');

 		// Restore code blocks
-		processed = processed.replace(/<<CODE_(\d+)>>/g, (_, index) => codeBlocks[parseInt(index)]);
+		processed = processed.replace(new RegExp(`${CODE_PLACEHOLDER_PREFIX}(\\d+)END`, 'g'), (_, index) => codeBlocks[parseInt(index)]);
+
+		// Clean up any remaining stray backslashes from unrecognized commands
+		processed = processed.replace(/\\(?=[a-zA-Z])/g, ''); // Remove \ before letters (unrecognized commands)

 		return processed;
 	}

 	/**
-	 * Render math expressions with KaTeX after HTML is generated
+	 * Render math expressions with KaTeX and restore HTML placeholders
 	 */
 	function renderMath(html: string): string {
-		// Render display math ($$...$$)
-		html = html.replace(/\$\$([\s\S]*?)\$\$/g, (_, math) => {
-			try {
-				return katex.renderToString(math.trim(), {
-					displayMode: true,
-					throwOnError: false,
-					output: 'html'
-				});
-			} catch {
-				return `<span class="math-error">$$${math}$$</span>`;
-			}
-		});
+		// Replace all math placeholders with rendered KaTeX
+		for (const [placeholder, { content, displayMode }] of mathExpressions) {
+			const escapedPlaceholder = placeholder.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+			const regex = new RegExp(escapedPlaceholder, 'g');

-		// Render inline math ($...$) but avoid matching currency like $5
-		html = html.replace(/\$([^\$\n]+?)\$/g, (match, math) => {
-			// Skip if it looks like currency ($ followed by number)
-			if (/^\d/.test(math.trim())) {
-				return match;
-			}
-			try {
-				return katex.renderToString(math.trim(), {
-					displayMode: false,
-					throwOnError: false,
-					output: 'html'
-				});
-			} catch {
-				return `<span class="math-error">$${math}$</span>`;
-			}
-		});
+			html = html.replace(regex, () => {
+				try {
+					const rendered = katex.renderToString(content, {
+						displayMode,
+						throwOnError: false,
+						output: 'html'
+					});
+
+					if (displayMode) {
+						return `
+							<div class="math-display-wrapper">
+								<div class="math-display-header">
+									<span class="math-label">LaTeX</span>
+									<button type="button" class="copy-math-btn" data-math-source="${encodeURIComponent(content)}" title="Copy LaTeX source">
+										<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+											<rect width="14" height="14" x="8" y="8" rx="2" ry="2"/>
+											<path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/>
+										</svg>
+									</button>
+								</div>
+								<div class="math-display-content">
+									${rendered}
+								</div>
+							</div>
+						`;
+					} else {
+						return `<span class="math-inline">${rendered}</span>`;
+					}
+				} catch {
+					const display = displayMode ? `$$${content}$$` : `$${content}$`;
+					return `<span class="math-error"><span class="math-error-icon">⚠</span> ${display}</span>`;
+				}
+			});
+		}
+
+		// Restore HTML placeholders (for \textbf, \emph, etc.)
+		for (const [placeholder, htmlContent] of htmlSnippets) {
+			const escapedPlaceholder = placeholder.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+			const regex = new RegExp(escapedPlaceholder, 'g');
+			html = html.replace(regex, htmlContent);
+		}

 		return html;
 	}
@@ -154,16 +377,50 @@
 		}
 	}

+	async function handleMathCopyClick(event: Event) {
+		const target = event.currentTarget as HTMLButtonElement;
+		const encodedSource = target.getAttribute('data-math-source');
+		if (!encodedSource) return;
+
+		const source = decodeURIComponent(encodedSource);
+
+		try {
+			await navigator.clipboard.writeText(source);
+			// Show copied feedback
+			const originalHtml = target.innerHTML;
+			target.innerHTML = `
+				<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+					<path d="M20 6L9 17l-5-5"/>
+				</svg>
+			`;
+			target.classList.add('copied');
+			setTimeout(() => {
+				target.innerHTML = originalHtml;
+				target.classList.remove('copied');
+			}, 2000);
+		} catch (error) {
+			console.error('Failed to copy math:', error);
+		}
+	}
+
 	function setupCopyButtons() {
 		if (!containerRef || !browser) return;

-		const buttons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
-		for (const button of buttons) {
+		const codeButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
+		for (const button of codeButtons) {
 			if (button.dataset.listenerBound !== 'true') {
 				button.dataset.listenerBound = 'true';
 				button.addEventListener('click', handleCopyClick);
 			}
 		}
+
+		const mathButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-math-btn');
+		for (const button of mathButtons) {
+			if (button.dataset.listenerBound !== 'true') {
+				button.dataset.listenerBound = 'true';
+				button.addEventListener('click', handleMathCopyClick);
+			}
+		}
 	}

 	$effect(() => {
@@ -424,28 +681,290 @@
 		color: #60a5fa;
 	}

-	/* KaTeX math styling */
+	/* KaTeX math styling - Base */
 	.markdown-content :global(.katex) {
 		font-size: 1.1em;
+		color: oklch(0.9 0 0);
 	}

-	.markdown-content :global(.katex-display) {
+	/* Display math container wrapper */
+	.markdown-content :global(.math-display-wrapper) {
 		margin: 1rem 0;
+		border-radius: 0.5rem;
+		overflow: hidden;
+		border: 1px solid rgba(255, 215, 0, 0.15);
+		background: rgba(0, 0, 0, 0.3);
+		transition: border-color 0.2s ease, box-shadow 0.2s ease;
+	}
+
+	.markdown-content :global(.math-display-wrapper:hover) {
+		border-color: rgba(255, 215, 0, 0.25);
+		box-shadow: 0 0 12px rgba(255, 215, 0, 0.08);
+	}
+
+	/* Display math header - hidden by default, slides in on hover */
+	.markdown-content :global(.math-display-header) {
+		display: flex;
+		justify-content: space-between;
+		align-items: center;
+		padding: 0.375rem 0.75rem;
+		background: rgba(255, 215, 0, 0.03);
+		border-bottom: 1px solid rgba(255, 215, 0, 0.08);
+		opacity: 0;
+		max-height: 0;
+		padding-top: 0;
+		padding-bottom: 0;
+		overflow: hidden;
+		transition:
+			opacity 0.2s ease,
+			max-height 0.2s ease,
+			padding 0.2s ease;
+	}
+
+	.markdown-content :global(.math-display-wrapper:hover .math-display-header) {
+		opacity: 1;
+		max-height: 2.5rem;
+		padding: 0.375rem 0.75rem;
+	}
+
+	.markdown-content :global(.math-label) {
+		color: rgba(255, 215, 0, 0.7);
+		font-size: 0.65rem;
+		font-weight: 500;
+		text-transform: uppercase;
+		letter-spacing: 0.1em;
+		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
+	}
+
+	.markdown-content :global(.copy-math-btn) {
+		display: flex;
+		align-items: center;
+		justify-content: center;
+		padding: 0.25rem;
+		background: transparent;
+		border: none;
+		color: var(--exo-light-gray, #9ca3af);
+		cursor: pointer;
+		transition: color 0.2s;
+		border-radius: 0.25rem;
+		opacity: 0;
+		transition:
+			color 0.2s,
+			opacity 0.15s ease;
+	}
+
+	.markdown-content :global(.math-display-wrapper:hover .copy-math-btn) {
+		opacity: 1;
+	}
+
+	.markdown-content :global(.copy-math-btn:hover) {
+		color: var(--exo-yellow, #ffd700);
+	}
+
+	.markdown-content :global(.copy-math-btn.copied) {
+		color: #22c55e;
+	}
+
+	/* Display math content area */
+	.markdown-content :global(.math-display-content) {
+		padding: 1rem 1.25rem;
 		overflow-x: auto;
 		overflow-y: hidden;
-		padding: 0.5rem 0;
 	}

-	.markdown-content :global(.katex-display > .katex) {
+	/* Custom scrollbar for math overflow */
+	.markdown-content :global(.math-display-content::-webkit-scrollbar) {
+		height: 6px;
+	}
+
+	.markdown-content :global(.math-display-content::-webkit-scrollbar-track) {
+		background: rgba(255, 255, 255, 0.05);
+		border-radius: 3px;
+	}
+
+	.markdown-content :global(.math-display-content::-webkit-scrollbar-thumb) {
+		background: rgba(255, 215, 0, 0.2);
+		border-radius: 3px;
+	}
+
+	.markdown-content :global(.math-display-content::-webkit-scrollbar-thumb:hover) {
+		background: rgba(255, 215, 0, 0.35);
+	}
+
+	.markdown-content :global(.math-display-content .katex-display) {
+		margin: 0;
+		padding: 0;
+	}
+
+	.markdown-content :global(.math-display-content .katex-display > .katex) {
 		text-align: center;
 	}

+	/* Inline math wrapper */
+	.markdown-content :global(.math-inline) {
+		display: inline;
+		padding: 0 0.125rem;
+		border-radius: 0.25rem;
+		transition: background-color 0.15s ease;
+	}
+
+	.markdown-content :global(.math-inline:hover) {
+		background: rgba(255, 215, 0, 0.05);
+	}
+
+	/* Dark theme KaTeX overrides */
+	.markdown-content :global(.katex .mord),
+	.markdown-content :global(.katex .minner),
+	.markdown-content :global(.katex .mop),
+	.markdown-content :global(.katex .mbin),
+	.markdown-content :global(.katex .mrel),
+	.markdown-content :global(.katex .mpunct) {
+		color: oklch(0.9 0 0);
+	}
+
+	/* Fraction lines and rules */
+	.markdown-content :global(.katex .frac-line),
+	.markdown-content :global(.katex .overline-line),
+	.markdown-content :global(.katex .underline-line),
+	.markdown-content :global(.katex .hline),
+	.markdown-content :global(.katex .rule) {
+		border-color: oklch(0.85 0 0) !important;
+		background: oklch(0.85 0 0);
+	}
+
+	/* Square roots and SVG elements */
+	.markdown-content :global(.katex .sqrt-line) {
+		border-color: oklch(0.85 0 0) !important;
+	}
+
+	.markdown-content :global(.katex svg) {
+		fill: oklch(0.85 0 0);
+		stroke: oklch(0.85 0 0);
+	}
+
+	.markdown-content :global(.katex svg path) {
+		stroke: oklch(0.85 0 0);
+	}
+
+	/* Delimiters (parentheses, brackets, braces) */
+	.markdown-content :global(.katex .delimsizing),
+	.markdown-content :global(.katex .delim-size1),
+	.markdown-content :global(.katex .delim-size2),
+	.markdown-content :global(.katex .delim-size3),
+	.markdown-content :global(.katex .delim-size4),
+	.markdown-content :global(.katex .mopen),
+	.markdown-content :global(.katex .mclose) {
+		color: oklch(0.75 0 0);
+	}
+
+	/* Math error styling */
 	.markdown-content :global(.math-error) {
+		display: inline-flex;
+		align-items: center;
+		gap: 0.375rem;
 		color: #f87171;
 		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
 		font-size: 0.875em;
 		background: rgba(248, 113, 113, 0.1);
-		padding: 0.125rem 0.25rem;
+		padding: 0.25rem 0.5rem;
 		border-radius: 0.25rem;
+		border: 1px solid rgba(248, 113, 113, 0.2);
+	}
+
+	.markdown-content :global(.math-error-icon) {
+		font-size: 0.875em;
+		opacity: 0.9;
+	}
+
+	/* LaTeX proof environment */
+	.markdown-content :global(.latex-proof) {
+		margin: 1rem 0;
+		padding: 1rem 1.25rem;
+		background: rgba(255, 255, 255, 0.02);
+		border-left: 3px solid rgba(255, 215, 0, 0.4);
+		border-radius: 0 0.375rem 0.375rem 0;
+	}
+
+	.markdown-content :global(.latex-proof-header) {
+		font-weight: 600;
+		font-style: italic;
+		color: oklch(0.85 0 0);
+		margin-bottom: 0.5rem;
+	}
+
+	.markdown-content :global(.latex-proof-header::after) {
+		content: '.';
+	}
+
+	.markdown-content :global(.latex-proof-content) {
+		color: oklch(0.9 0 0);
+	}
+
+	.markdown-content :global(.latex-proof-content p:last-child) {
+		margin-bottom: 0;
+	}
+
+	/* QED symbol at end of proof */
+	.markdown-content :global(.latex-proof-content::after) {
+		content: '∎';
+		display: block;
+		text-align: right;
+		color: oklch(0.7 0 0);
+		margin-top: 0.5rem;
+	}
+
+	/* LaTeX theorem-like environments */
+	.markdown-content :global(.latex-theorem) {
+		margin: 1rem 0;
+		padding: 1rem 1.25rem;
+		background: rgba(255, 215, 0, 0.03);
+		border: 1px solid rgba(255, 215, 0, 0.15);
+		border-radius: 0.375rem;
+	}
+
+	.markdown-content :global(.latex-theorem-header) {
+		font-weight: 700;
+		color: var(--exo-yellow, #ffd700);
+		margin-bottom: 0.5rem;
+	}
+
+	.markdown-content :global(.latex-theorem-header::after) {
+		content: '.';
+	}
+
+	.markdown-content :global(.latex-theorem-content) {
+		color: oklch(0.9 0 0);
+		font-style: italic;
+	}
+
+	.markdown-content :global(.latex-theorem-content p:last-child) {
+		margin-bottom: 0;
+	}
+
+	/* LaTeX diagram/figure placeholder */
+	.markdown-content :global(.latex-diagram-placeholder) {
+		display: flex;
+		align-items: center;
+		justify-content: center;
+		gap: 0.5rem;
+		margin: 1rem 0;
+		padding: 1.5rem 2rem;
+		background: rgba(255, 255, 255, 0.02);
+		border: 1px dashed rgba(255, 215, 0, 0.25);
+		border-radius: 0.5rem;
+		color: rgba(255, 215, 0, 0.6);
+		font-size: 0.875rem;
+	}
+
+	.markdown-content :global(.latex-diagram-icon) {
+		font-size: 1.25rem;
+		opacity: 0.8;
+	}
+
+	.markdown-content :global(.latex-diagram-text) {
+		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
+		font-size: 0.75rem;
+		text-transform: uppercase;
+		letter-spacing: 0.05em;
 	}
 </style>
--- a/dashboard/src/lib/components/ModelCard.svelte
+++ b/dashboard/src/lib/components/ModelCard.svelte
@@ -197,7 +197,7 @@ function toggleNodeDetails(nodeId: string): void {
 	// Uses API preview data when available, falls back to local estimation
 	const placementPreview = $derived(() => {
 		const nodeArray = nodeList();
-		if (nodeArray.length === 0) return { nodes: [], canFit: false, totalAvailable: 0, error: null };
+		if (nodeArray.length === 0) return { nodes: [], canFit: false, totalAvailable: 0, topoWidth: 260, topoHeight: 90, error: null };
 		
 		const numNodes = nodeArray.length;
 		const iconSize = numNodes === 1 ? 50 : 36;
--- a/dashboard/src/lib/components/TokenHeatmap.svelte
+++ b/dashboard/src/lib/components/TokenHeatmap.svelte
@@ -1,192 +0,0 @@
-<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>
--- a/dashboard/src/lib/components/TopologyGraph.svelte
+++ b/dashboard/src/lib/components/TopologyGraph.svelte
@@ -1,7 +1,7 @@
 <script lang="ts">
 	import { onMount, onDestroy } from 'svelte';
 	import * as d3 from 'd3';
-import { topologyData, isTopologyMinimized, debugMode } from '$lib/stores/app.svelte';
+import { topologyData, isTopologyMinimized, debugMode, type NodeInfo } from '$lib/stores/app.svelte';

 	interface Props {
 		class?: string;
@@ -24,14 +24,14 @@ function getNodeLabel(nodeId: string): string {

 function getInterfaceLabel(nodeId: string, ip?: string): { label: string; missing: boolean } {
 	if (!ip) return { label: '?', missing: true };
-	
+
 	// Strip port if present (e.g., "192.168.1.1:8080" -> "192.168.1.1")
 	const cleanIp = ip.includes(':') && !ip.includes('[') ? ip.split(':')[0] : ip;
-	
+
 	// Helper to check a node's interfaces
-	function checkNode(node: typeof data.nodes[string]): string | null {
+	function checkNode(node: NodeInfo | undefined): string | null {
 		if (!node) return null;
-		
+
 		const matchFromInterfaces = node.network_interfaces?.find((iface) =>
 			(iface.addresses || []).some((addr) => addr === cleanIp || addr === ip)
 		);
@@ -39,17 +39,19 @@ function getInterfaceLabel(nodeId: string, ip?: string): { label: string; missin
 			return matchFromInterfaces.name;
 		}

-		const mapped = node.ip_to_interface?.[cleanIp] || node.ip_to_interface?.[ip];
-		if (mapped && mapped.trim().length > 0) {
-			return mapped;
+		if (node.ip_to_interface) {
+			const mapped = node.ip_to_interface[cleanIp] || (ip ? node.ip_to_interface[ip] : undefined);
+			if (mapped && mapped.trim().length > 0) {
+				return mapped;
+			}
 		}
 		return null;
 	}
-	
+
 	// Try specified node first
 	const result = checkNode(data?.nodes?.[nodeId]);
 	if (result) return { label: result, missing: false };
-	
+
 	// Fallback: search all nodes for this IP
 	for (const [, otherNode] of Object.entries(data?.nodes || {})) {
 		const otherResult = checkNode(otherNode);
@@ -255,21 +257,24 @@ function wrapLine(text: string, maxLen: number): string[] {
 		const arrowsGroup = svg.append('g').attr('class', 'arrows-group');
 		const debugLabelsGroup = svg.append('g').attr('class', 'debug-edge-labels');

-		const pairMap = new Map<string, { a: string; b: string; aToB: boolean; bToA: boolean; connections: Array<{ from: string; to: string; ip: string; ifaceLabel: string; missingIface: boolean }> }>();
-		let debugEdgeLabels: Array<{ connections: typeof pairMap extends Map<string, infer V> ? V['connections'] : never; isLeft: boolean; isTop: boolean; mx: number; my: number }> | null = null;
+		type ConnectionInfo = { from: string; to: string; ip: string; ifaceLabel: string; missingIface: boolean };
+		type PairEntry = { a: string; b: string; aToB: boolean; bToA: boolean; connections: ConnectionInfo[] };
+		type DebugEdgeLabelEntry = { connections: ConnectionInfo[]; isLeft: boolean; isTop: boolean; mx: number; my: number };
+		const pairMap = new Map<string, PairEntry>();
+		const debugEdgeLabels: DebugEdgeLabelEntry[] = [];
 		edges.forEach(edge => {
 			if (!edge.source || !edge.target || edge.source === edge.target) return;
 			if (!positionById[edge.source] || !positionById[edge.target]) return;
-			
+
 			const a = edge.source < edge.target ? edge.source : edge.target;
 			const b = edge.source < edge.target ? edge.target : edge.source;
 			const key = `${a}|${b}`;
 			const entry = pairMap.get(key) || { a, b, aToB: false, bToA: false, connections: [] };
-			
+
 			if (edge.source === a) entry.aToB = true;
 			else entry.bToA = true;

-			const ip = edge.sendBackIp || edge.sendBackMultiaddr?.ip_address || '?';
+			const ip = edge.sendBackIp || '?';
 			const ifaceInfo = getInterfaceLabel(edge.source, ip);
 			entry.connections.push({
 				from: edge.source,
@@ -338,9 +343,8 @@ function wrapLine(text: string, maxLen: number): string[] {
 				// Determine which side of viewport based on edge midpoint
 				const isLeft = mx < centerX;
 				const isTop = my < safeCenterY;
-				
+
 				// Store for batch rendering after all edges processed
-				if (!debugEdgeLabels) debugEdgeLabels = [];
 				debugEdgeLabels.push({
 					connections: entry.connections,
 					isLeft,
@@ -381,32 +385,32 @@ function wrapLine(text: string, maxLen: number): string[] {
 			}
 			
 			// Group by quadrant: topLeft, topRight, bottomLeft, bottomRight
-			const quadrants: Record<string, typeof debugEdgeLabels> = {
+			const quadrants: Record<string, DebugEdgeLabelEntry[]> = {
 				topLeft: [],
 				topRight: [],
 				bottomLeft: [],
 				bottomRight: []
 			};
-			
+
 			debugEdgeLabels.forEach(edge => {
 				const key = (edge.isTop ? 'top' : 'bottom') + (edge.isLeft ? 'Left' : 'Right');
 				quadrants[key].push(edge);
 			});
-			
+
 			// Render each quadrant
-			Object.entries(quadrants).forEach(([quadrant, edges]) => {
-				if (edges.length === 0) return;
-				
+			Object.entries(quadrants).forEach(([quadrant, quadrantEdges]) => {
+				if (quadrantEdges.length === 0) return;
+
 				const isLeft = quadrant.includes('Left');
 				const isTop = quadrant.includes('top');
-				
+
 				let baseX = isLeft ? padding : width - padding;
 				let baseY = isTop ? padding : height - padding;
 				const textAnchor = isLeft ? 'start' : 'end';
-				
+
 				let currentY = baseY;
-				
-				edges.forEach(edge => {
+
+				quadrantEdges.forEach(edge => {
 					edge.connections.forEach(conn => {
 						const arrow = getArrow(conn.from, conn.to);
 						const label = `${arrow} ${conn.ip} ${conn.ifaceLabel}`;
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -71,52 +71,66 @@ export interface Instance {
 	};
 }

-interface RawNodeProfile {
+// Granular node state types from the new state structure
+interface RawNodeIdentity {
 	modelId?: string;
 	chipId?: string;
 	friendlyName?: string;
-	networkInterfaces?: Array<{
-		name?: string;
-		ipAddress?: string;
-		addresses?: Array<{ address?: string } | string>;
-		ipv4?: string;
-		ipv6?: string;
-		ipAddresses?: string[];
-		ips?: string[];
-	}>;
-	memory?: {
-		ramTotal?: { inBytes: number };
-		ramAvailable?: { inBytes: number };
-		swapTotal?: { inBytes: number };
-		swapAvailable?: { inBytes: number };
-	};
-	system?: {
-		gpuUsage?: number;
-		temp?: number;
-		sysPower?: number;
+}
+
+interface RawMemoryUsage {
+	ramTotal?: { inBytes: number };
+	ramAvailable?: { inBytes: number };
+	swapTotal?: { inBytes: number };
+	swapAvailable?: { inBytes: number };
+}
+
+interface RawSystemPerformanceProfile {
+	gpuUsage?: number;
+	temp?: number;
+	sysPower?: number;
+	pcpuUsage?: number;
+	ecpuUsage?: number;
+}
+
+interface RawNetworkInterfaceInfo {
+	name?: string;
+	ipAddress?: string;
+	addresses?: Array<{ address?: string } | string>;
+	ipv4?: string;
+	ipv6?: string;
+	ipAddresses?: string[];
+	ips?: string[];
+}
+
+interface RawNodeNetworkInfo {
+	interfaces?: RawNetworkInterfaceInfo[];
+}
+
+interface RawSocketConnection {
+	sinkMultiaddr?: {
+		address?: string;
+		ip_address?: string;
+		address_type?: string;
+		port?: number;
 	};
 }

-interface RawTopologyNode {
-	nodeId: string;
-	nodeProfile: RawNodeProfile;
+interface RawRDMAConnection {
+	sourceRdmaIface?: string;
+	sinkRdmaIface?: string;
 }

-interface RawTopologyConnection {
-	localNodeId: string;
-	sendBackNodeId: string;
-	sendBackMultiaddr?:
-		| { multiaddr?: string; address?: string; ip_address?: string }
-		| string;
-}
+type RawConnectionEdge = RawSocketConnection | RawRDMAConnection;
+
+// New nested mapping format: { source: { sink: [edge1, edge2, ...] } }
+type RawConnectionsMap = Record<string, Record<string, RawConnectionEdge[]>>;

 interface RawTopology {
-	nodes: RawTopologyNode[];
-	connections?: RawTopologyConnection[];
+	nodes: string[];
+	connections?: RawConnectionsMap;
 }

-type RawNodeProfiles = Record<string, RawNodeProfile>;
-
 export interface DownloadProgress {
 	totalBytes: number;
 	downloadedBytes: number;
@@ -171,7 +185,11 @@ interface RawStateResponse {
 	>;
 	runners?: Record<string, unknown>;
 	downloads?: Record<string, unknown[]>;
-	nodeProfiles?: RawNodeProfiles;
+	// New granular node state fields
+	nodeIdentities?: Record<string, RawNodeIdentity>;
+	nodeMemory?: Record<string, RawMemoryUsage>;
+	nodeSystem?: Record<string, RawSystemPerformanceProfile>;
+	nodeNetwork?: Record<string, RawNodeNetworkInfo>;
 }

 export interface MessageAttachment {
@@ -182,20 +200,6 @@ 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";
@@ -205,7 +209,6 @@ 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 {
@@ -221,56 +224,69 @@ export interface Conversation {

 const STORAGE_KEY = "exo-conversations";

+interface GranularNodeState {
+	nodeIdentities?: Record<string, RawNodeIdentity>;
+	nodeMemory?: Record<string, RawMemoryUsage>;
+	nodeSystem?: Record<string, RawSystemPerformanceProfile>;
+	nodeNetwork?: Record<string, RawNodeNetworkInfo>;
+}
+
+function transformNetworkInterface(iface: RawNetworkInterfaceInfo): {
+	name?: string;
+	addresses: string[];
+} {
+	const addresses: string[] = [];
+	if (iface.ipAddress && typeof iface.ipAddress === "string") {
+		addresses.push(iface.ipAddress);
+	}
+	if (Array.isArray(iface.addresses)) {
+		for (const addr of iface.addresses) {
+			if (typeof addr === "string") addresses.push(addr);
+			else if (addr && typeof addr === "object" && addr.address)
+				addresses.push(addr.address);
+		}
+	}
+	if (Array.isArray(iface.ipAddresses)) {
+		addresses.push(
+			...iface.ipAddresses.filter((a): a is string => typeof a === "string"),
+		);
+	}
+	if (Array.isArray(iface.ips)) {
+		addresses.push(
+			...iface.ips.filter((a): a is string => typeof a === "string"),
+		);
+	}
+	if (iface.ipv4 && typeof iface.ipv4 === "string") addresses.push(iface.ipv4);
+	if (iface.ipv6 && typeof iface.ipv6 === "string") addresses.push(iface.ipv6);
+
+	return {
+		name: iface.name,
+		addresses: Array.from(new Set(addresses)),
+	};
+}
+
 function transformTopology(
 	raw: RawTopology,
-	profiles?: RawNodeProfiles,
+	granularState: GranularNodeState,
 ): TopologyData {
 	const nodes: Record<string, NodeInfo> = {};
 	const edges: TopologyEdge[] = [];

-	for (const node of raw.nodes || []) {
-		const mergedProfile = profiles?.[node.nodeId];
-		const profile = { ...(node.nodeProfile ?? {}), ...(mergedProfile ?? {}) };
-		const ramTotal = profile?.memory?.ramTotal?.inBytes ?? 0;
-		const ramAvailable = profile?.memory?.ramAvailable?.inBytes ?? 0;
+	for (const nodeId of raw.nodes || []) {
+		if (!nodeId) continue;
+
+		// Get data from granular state mappings
+		const identity = granularState.nodeIdentities?.[nodeId];
+		const memory = granularState.nodeMemory?.[nodeId];
+		const system = granularState.nodeSystem?.[nodeId];
+		const network = granularState.nodeNetwork?.[nodeId];
+
+		const ramTotal = memory?.ramTotal?.inBytes ?? 0;
+		const ramAvailable = memory?.ramAvailable?.inBytes ?? 0;
 		const ramUsage = Math.max(ramTotal - ramAvailable, 0);

-		const networkInterfaces = (profile?.networkInterfaces || []).map(
-			(iface) => {
-				const addresses: string[] = [];
-				if (iface.ipAddress && typeof iface.ipAddress === "string") {
-					addresses.push(iface.ipAddress);
-				}
-				if (Array.isArray(iface.addresses)) {
-					for (const addr of iface.addresses) {
-						if (typeof addr === "string") addresses.push(addr);
-						else if (addr && typeof addr === "object" && addr.address)
-							addresses.push(addr.address);
-					}
-				}
-				if (Array.isArray(iface.ipAddresses)) {
-					addresses.push(
-						...iface.ipAddresses.filter(
-							(a): a is string => typeof a === "string",
-						),
-					);
-				}
-				if (Array.isArray(iface.ips)) {
-					addresses.push(
-						...iface.ips.filter((a): a is string => typeof a === "string"),
-					);
-				}
-				if (iface.ipv4 && typeof iface.ipv4 === "string")
-					addresses.push(iface.ipv4);
-				if (iface.ipv6 && typeof iface.ipv6 === "string")
-					addresses.push(iface.ipv6);
-
-				return {
-					name: iface.name,
-					addresses: Array.from(new Set(addresses)),
-				};
-			},
-		);
+		const rawInterfaces = network?.interfaces || [];
+		const networkInterfaces = rawInterfaces.map(transformNetworkInterface);

 		const ipToInterface: Record<string, string> = {};
 		for (const iface of networkInterfaces) {
@@ -279,10 +295,10 @@ function transformTopology(
 			}
 		}

-		nodes[node.nodeId] = {
+		nodes[nodeId] = {
 			system_info: {
-				model_id: profile?.modelId ?? "Unknown",
-				chip: profile?.chipId,
+				model_id: identity?.modelId ?? "Unknown",
+				chip: identity?.chipId,
 				memory: ramTotal,
 			},
 			network_interfaces: networkInterfaces,
@@ -293,43 +309,42 @@ function transformTopology(
 					ram_total: ramTotal,
 				},
 				temp:
-					profile?.system?.temp !== undefined
-						? { gpu_temp_avg: profile.system.temp }
+					system?.temp !== undefined
+						? { gpu_temp_avg: system.temp }
 						: undefined,
 				gpu_usage:
-					profile?.system?.gpuUsage !== undefined
-						? [0, profile.system.gpuUsage]
-						: undefined,
-				sys_power: profile?.system?.sysPower,
+					system?.gpuUsage !== undefined ? [0, system.gpuUsage] : undefined,
+				sys_power: system?.sysPower,
 			},
 			last_macmon_update: Date.now() / 1000,
-			friendly_name: profile?.friendlyName,
+			friendly_name: identity?.friendlyName,
 		};
 	}

-	for (const conn of raw.connections || []) {
-		if (!conn.localNodeId || !conn.sendBackNodeId) continue;
-		if (conn.localNodeId === conn.sendBackNodeId) continue;
-		if (!nodes[conn.localNodeId] || !nodes[conn.sendBackNodeId]) continue;
+	// Handle connections - nested mapping format { source: { sink: [edges] } }
+	const connections = raw.connections;
+	if (connections && typeof connections === "object") {
+		for (const [source, sinks] of Object.entries(connections)) {
+			if (!sinks || typeof sinks !== "object") continue;
+			for (const [sink, edgeList] of Object.entries(sinks)) {
+				if (!Array.isArray(edgeList)) continue;
+				for (const edge of edgeList) {
+					let sendBackIp: string | undefined;
+					if (edge && typeof edge === "object" && "sinkMultiaddr" in edge) {
+						const multiaddr = edge.sinkMultiaddr;
+						if (multiaddr) {
+							sendBackIp =
+								multiaddr.ip_address ||
+								extractIpFromMultiaddr(multiaddr.address);
+						}
+					}

-		let sendBackIp: string | undefined;
-		if (conn.sendBackMultiaddr) {
-			const multi = conn.sendBackMultiaddr;
-			if (typeof multi === "string") {
-				sendBackIp = extractIpFromMultiaddr(multi);
-			} else {
-				sendBackIp =
-					multi.ip_address ||
-					extractIpFromMultiaddr(multi.multiaddr) ||
-					extractIpFromMultiaddr(multi.address);
+					if (nodes[source] && nodes[sink] && source !== sink) {
+						edges.push({ source, target: sink, sendBackIp });
+					}
+				}
 			}
 		}
-
-		edges.push({
-			source: conn.localNodeId,
-			target: conn.sendBackNodeId,
-			sendBackIp,
-		});
 	}

 	return { nodes, edges };
@@ -383,21 +398,6 @@ 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,61 +435,12 @@ 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 {
-			// 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));
+			localStorage.setItem(STORAGE_KEY, JSON.stringify(this.conversations));
 		} 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");
 		}
 	}

@@ -947,7 +898,12 @@ class AppStore {
 			const data: RawStateResponse = await response.json();

 			if (data.topology) {
-				this.topologyData = transformTopology(data.topology, data.nodeProfiles);
+				this.topologyData = transformTopology(data.topology, {
+					nodeIdentities: data.nodeIdentities,
+					nodeMemory: data.nodeMemory,
+					nodeSystem: data.nodeSystem,
+					nodeNetwork: data.nodeNetwork,
+				});
 			}
 			if (data.instances) {
 				this.instances = data.instances;
@@ -1410,11 +1366,6 @@ 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 = {
@@ -1492,10 +1443,7 @@ class AppStore {
 					messages: apiMessages,
 					temperature: 0.7,
 					stream: true,
-					logprobs: true,
-					top_logprobs: 5,
 				}),
-				signal: controller.signal,
 			});

 			if (!response.ok) {
@@ -1511,7 +1459,6 @@ class AppStore {
 			const decoder = new TextDecoder();
 			let fullContent = "";
 			let buffer = "";
-			const collectedTokens: TokenData[] = [];

 			while (true) {
 				const { done, value } = await reader.read();
@@ -1533,8 +1480,8 @@ class AppStore {

 						try {
 							const parsed = JSON.parse(data);
-							const delta = parsed.choices?.[0]?.delta?.content;
-							if (delta) {
+							const tokenContent = parsed.choices?.[0]?.delta?.content;
+							if (tokenContent) {
 								// Track first token for TTFT
 								if (firstTokenTime === null) {
 									firstTokenTime = performance.now();
@@ -1551,30 +1498,7 @@ class AppStore {
 									this.tps = (tokenCount / elapsed) * 1000;
 								}

-								// 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;
+								fullContent += tokenContent;

 								// Strip thinking tags for display and extract thinking content
 								const { displayContent, thinkingContent } =
@@ -1588,7 +1512,6 @@ class AppStore {
 								if (idx !== -1) {
 									this.messages[idx].content = displayContent;
 									this.messages[idx].thinking = thinkingContent || undefined;
-									this.messages[idx].tokens = [...collectedTokens];
 								}
 								this.persistActiveConversation();
 							}
@@ -1636,16 +1559,9 @@ 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);
@@ -1655,237 +1571,6 @@ 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();
@@ -1965,8 +1650,6 @@ 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;
--- a/dashboard/src/routes/+page.svelte
+++ b/dashboard/src/routes/+page.svelte
@@ -915,7 +915,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 		const runnerEntries = Object.entries(runnerToShard).map(([runnerId, shardWrapped]) => {
 			const [tag, shard] = getTagged(shardWrapped);
 			const meta = (shard as { modelMeta?: { worldSize?: number; nLayers?: number; deviceRank?: number } } | undefined);
-			const deviceRank = (meta?.deviceRank as number | undefined) ?? 0;
+			const deviceRank = meta?.modelMeta?.deviceRank ?? 0;
 			return { runnerId, tag, deviceRank };
 		});

--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,7 @@ dependencies = [
 exo-master = "exo.master.main:main"
 exo-worker = "exo.worker.main:main"
 exo = "exo.main:main"
+exo-rsh = "exo.rsh.client:main"

 # dependencies only required for development
 [dependency-groups]
--- a/src/exo/cli/init.py
+++ b/src/exo/cli/init.py
@@ -0,0 +1,32 @@
+"""Exo CLI - SLURM-compatible job management commands."""
+
+
+def run_subcommand(command: str, args: list[str]) -> int:
+    """Route to the appropriate subcommand handler.
+
+    Args:
+        command: The subcommand name (sbatch, squeue, scancel, salloc)
+        args: Command line arguments for the subcommand
+
+    Returns:
+        Exit code from the subcommand
+    """
+    if command == "sbatch":
+        from exo.cli.sbatch import main
+
+        return main(args)
+    elif command == "squeue":
+        from exo.cli.squeue import main
+
+        return main(args)
+    elif command == "scancel":
+        from exo.cli.scancel import main
+
+        return main(args)
+    elif command == "salloc":
+        from exo.cli.salloc import main
+
+        return main(args)
+    else:
+        print(f"Unknown subcommand: {command}")
+        return 1
--- a/src/exo/cli/common.py
+++ b/src/exo/cli/common.py
@@ -0,0 +1,118 @@
+"""Common utilities for Exo CLI commands."""
+
+import json
+import os
+import urllib.request
+from typing import Any
+from urllib.error import HTTPError, URLError
+
+# Default API endpoint
+DEFAULT_API_HOST = "localhost"
+DEFAULT_API_PORT = 52415
+
+
+def get_api_base() -> str:
+    """Get the API base URL from environment or defaults."""
+    host = os.environ.get("EXO_API_HOST", DEFAULT_API_HOST)
+    port = os.environ.get("EXO_API_PORT", str(DEFAULT_API_PORT))
+    return f"http://{host}:{port}"
+
+
+def api_request(
+    method: str,
+    path: str,
+    data: dict[str, Any] | None = None,
+) -> dict[str, Any] | list[Any]:
+    """Make an API request to the Exo server.
+
+    Args:
+        method: HTTP method (GET, POST, DELETE, etc.)
+        path: API path (e.g., "/flash/instances")
+        data: Optional JSON data for POST/PUT requests
+
+    Returns:
+        Parsed JSON response
+
+    Raises:
+        SystemExit: On connection or HTTP errors
+    """
+    url = f"{get_api_base()}{path}"
+
+    request_data = None
+    if data is not None:
+        request_data = json.dumps(data).encode("utf-8")
+
+    req = urllib.request.Request(
+        url,
+        data=request_data,
+        method=method,
+    )
+    req.add_header("Content-Type", "application/json")
+
+    try:
+        with urllib.request.urlopen(req, timeout=30) as response:  # pyright: ignore[reportAny]
+            body: str = response.read().decode("utf-8")  # pyright: ignore[reportAny]
+            if body:
+                return json.loads(body)  # pyright: ignore[reportAny]
+            return {}
+    except HTTPError as e:
+        error_body = e.read().decode("utf-8") if e.fp else ""
+        print(f"API error: {e.code} {e.reason}")
+        if error_body:
+            try:
+                error_json: dict[str, str] = json.loads(error_body)  # pyright: ignore[reportAny]
+                if "detail" in error_json:
+                    print(f"  {error_json['detail']}")
+            except json.JSONDecodeError:
+                print(f"  {error_body}")
+        raise SystemExit(1) from None
+    except URLError as e:
+        print(f"Connection error: {e.reason}")
+        print(f"Is Exo running at {get_api_base()}?")
+        raise SystemExit(1) from None
+
+
+def truncate_id(instance_id: str, length: int = 8) -> str:
+    """Truncate a UUID for display.
+
+    Args:
+        instance_id: Full UUID string
+        length: Number of characters to keep
+
+    Returns:
+        Truncated ID without hyphens
+    """
+    return instance_id.replace("-", "")[:length]
+
+
+def format_table(headers: list[str], rows: list[list[str]]) -> str:
+    """Format data as a simple text table.
+
+    Args:
+        headers: Column headers
+        rows: List of rows, each row is a list of column values
+
+    Returns:
+        Formatted table string
+    """
+    if not rows:
+        return "  ".join(f"{h:<10}" for h in headers)
+
+    # Calculate column widths
+    widths = [len(h) for h in headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            if i < len(widths):
+                widths[i] = max(widths[i], len(cell))
+
+    # Build format string
+    fmt = "  ".join(f"{{:<{w}}}" for w in widths)
+
+    # Format output
+    lines = [fmt.format(*headers)]
+    for row in rows:
+        # Pad row if needed
+        padded = row + [""] * (len(headers) - len(row))
+        lines.append(fmt.format(*padded[: len(headers)]))
+
+    return "\n".join(lines)
--- a/src/exo/cli/salloc.py
+++ b/src/exo/cli/salloc.py
@@ -0,0 +1,100 @@
+"""salloc - Allocate nodes for interactive use.
+
+Usage:
+    exo salloc [options] [-- command [args...]]
+
+Options:
+    -N, --nodes N       Number of nodes to allocate (default: 1)
+    --hosts HOSTS       Comma-separated list of hostnames
+
+If a command is provided after --, it will be executed with
+SLURM-like environment variables set:
+    SLURM_JOB_NODELIST  - Comma-separated list of allocated nodes
+    SLURM_NNODES        - Number of allocated nodes
+
+Examples:
+    exo salloc --nodes=2 --hosts=node1,node2 -- mpirun ./my_program
+    exo salloc --hosts=localhost -- bash
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for salloc command."""
+    # Split args at -- if present
+    cmd_args: list[str] = []
+    salloc_args = args
+
+    if "--" in args:
+        idx = args.index("--")
+        salloc_args = args[:idx]
+        cmd_args = args[idx + 1 :]
+
+    parser = argparse.ArgumentParser(
+        prog="exo salloc",
+        description="Allocate nodes for interactive use",
+    )
+    parser.add_argument(
+        "-N",
+        "--nodes",
+        type=int,
+        default=1,
+        help="Number of nodes to allocate (default: 1)",
+    )
+    parser.add_argument(
+        "--hosts",
+        help="Comma-separated list of hostnames (required)",
+    )
+
+    parsed = parser.parse_args(salloc_args)
+
+    nodes: int = parsed.nodes  # pyright: ignore[reportAny]
+    hosts: str | None = parsed.hosts  # pyright: ignore[reportAny]
+
+    # Require explicit hosts since we can't discover them from topology
+    if not hosts:
+        print("Error: --hosts is required (e.g., --hosts=node1,node2)", file=sys.stderr)
+        print("       The Exo topology doesn't expose hostnames.", file=sys.stderr)
+        return 1
+
+    host_list = [h.strip() for h in hosts.split(",") if h.strip()]
+
+    if len(host_list) < nodes:
+        print(
+            f"Error: Requested {nodes} nodes but only {len(host_list)} hosts provided",
+            file=sys.stderr,
+        )
+        return 1
+
+    # Use first N hosts
+    allocated_hosts = host_list[:nodes]
+    nodelist = ",".join(allocated_hosts)
+
+    # Set environment variables
+    env = os.environ.copy()
+    env["SLURM_JOB_NODELIST"] = nodelist
+    env["SLURM_NNODES"] = str(nodes)
+
+    print(f"salloc: Granted job allocation on {nodes} node(s)")
+    print(f"salloc: Nodes: {nodelist}")
+
+    if cmd_args:
+        # Run the command
+        print(f"salloc: Running: {' '.join(cmd_args)}")
+        result = subprocess.run(cmd_args, env=env)
+        return result.returncode
+    else:
+        # Start interactive shell
+        shell = os.environ.get("SHELL", "/bin/bash")
+        print(f"salloc: Starting shell {shell}")
+        print("salloc: Use 'exit' to release allocation")
+        result = subprocess.run([shell], env=env)
+        return result.returncode
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
--- a/src/exo/cli/sbatch.py
+++ b/src/exo/cli/sbatch.py
@@ -0,0 +1,233 @@
+"""sbatch - Submit a batch job to Exo.
+
+Usage:
+    exo sbatch [options] <script|executable>
+    exo sbatch --job-name=NAME --nodes=N <executable>
+
+Options:
+    -J, --job-name NAME         Job name
+    -N, --nodes N               Number of nodes (default: 1)
+    --ntasks-per-node N         Tasks per node (default: 1)
+    -D, --chdir DIR             Working directory
+    --hosts HOSTS               Comma-separated list of hostnames
+
+Job scripts can contain #SBATCH directives:
+    #!/bin/bash
+    #SBATCH --job-name=Sod2D
+    #SBATCH --nodes=2
+    #SBATCH --chdir=/path/to/workdir
+
+    /path/to/flash4
+"""
+
+import argparse
+import os
+import re
+import sys
+
+from exo.cli.common import api_request, truncate_id
+
+
+def parse_job_script(script_path: str) -> tuple[dict[str, str], str | None]:
+    """Parse a job script for #SBATCH directives and executable.
+
+    Args:
+        script_path: Path to the job script
+
+    Returns:
+        Tuple of (directives dict, executable path or None)
+    """
+    directives: dict[str, str] = {}
+    executable: str | None = None
+
+    with open(script_path, "r") as f:
+        for line in f:
+            line = line.strip()
+
+            # Parse #SBATCH directives
+            if line.startswith("#SBATCH"):
+                # Handle both --option=value and --option value formats
+                match = re.match(r"#SBATCH\s+(-\w|--[\w-]+)(?:=|\s+)(.+)", line)
+                if match:
+                    opt, val = match.groups()
+                    directives[opt.lstrip("-")] = val.strip()
+                continue
+
+            # Skip comments and empty lines
+            if line.startswith("#") or not line:
+                continue
+
+            # First non-comment, non-directive line is the executable
+            if executable is None:
+                # Handle lines like "/path/to/flash4" or "srun /path/to/flash4"
+                parts = line.split()
+                if parts:
+                    # Skip srun/mpirun prefixes if present
+                    for part in parts:
+                        if not part.startswith("-") and "/" in part:
+                            executable = part
+                            break
+                    if executable is None and parts:
+                        executable = parts[-1]  # Last token
+
+    return directives, executable
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for sbatch command."""
+    parser = argparse.ArgumentParser(
+        prog="exo sbatch",
+        description="Submit a batch job to Exo",
+    )
+    parser.add_argument(
+        "script",
+        help="Job script or executable path",
+    )
+    parser.add_argument(
+        "-J",
+        "--job-name",
+        dest="job_name",
+        help="Job name",
+    )
+    parser.add_argument(
+        "-N",
+        "--nodes",
+        type=int,
+        default=1,
+        help="Number of nodes (default: 1)",
+    )
+    parser.add_argument(
+        "--ntasks-per-node",
+        type=int,
+        default=1,
+        help="Tasks per node (default: 1)",
+    )
+    parser.add_argument(
+        "-D",
+        "--chdir",
+        help="Working directory",
+    )
+    parser.add_argument(
+        "--hosts",
+        help="Comma-separated list of hostnames",
+    )
+
+    parsed = parser.parse_args(args)
+
+    # Extract typed values from namespace
+    script_path: str = parsed.script  # pyright: ignore[reportAny]
+    arg_job_name: str | None = parsed.job_name  # pyright: ignore[reportAny]
+    arg_nodes: int = parsed.nodes  # pyright: ignore[reportAny]
+    arg_ntasks: int = parsed.ntasks_per_node  # pyright: ignore[reportAny]
+    arg_chdir: str | None = parsed.chdir  # pyright: ignore[reportAny]
+    arg_hosts: str | None = parsed.hosts  # pyright: ignore[reportAny]
+
+    # Determine if input is a script or direct executable
+    executable: str | None = None
+    directives: dict[str, str] = {}
+
+    if os.path.isfile(script_path):
+        # Check if it's a binary file (executable) or text script
+        is_binary = False
+        try:
+            with open(script_path, "rb") as f:
+                chunk = f.read(512)
+                # Binary files typically contain null bytes
+                is_binary = b"\x00" in chunk
+        except OSError:
+            pass
+
+        if is_binary:
+            # It's a binary executable
+            executable = script_path
+        else:
+            # Try to read as text
+            try:
+                with open(script_path, "r") as f:
+                    first_line = f.readline()
+                    f.seek(0)
+                    content = f.read(1024)
+
+                if first_line.startswith("#!") or "#SBATCH" in content:
+                    # It's a job script - parse it
+                    directives, executable = parse_job_script(script_path)
+                else:
+                    # It's an executable (text but no shebang/directives)
+                    executable = script_path
+            except UnicodeDecodeError:
+                # Can't read as text - treat as binary executable
+                executable = script_path
+    else:
+        # Not a file - treat as executable path
+        executable = script_path
+
+    if executable is None:
+        print("Error: No executable found in job script", file=sys.stderr)
+        return 1
+
+    # Build job parameters - CLI args override script directives
+    job_name = arg_job_name or directives.get("job-name") or directives.get("J")
+    if not job_name:
+        # Generate name from executable
+        job_name = os.path.basename(executable).replace(".", "_")
+
+    nodes = arg_nodes
+    if "nodes" in directives:
+        nodes = int(directives["nodes"])
+    if "N" in directives:
+        nodes = int(directives["N"])
+    if arg_nodes != 1:  # CLI override
+        nodes = arg_nodes
+
+    ntasks = arg_ntasks
+    if "ntasks-per-node" in directives:
+        ntasks = int(directives["ntasks-per-node"])
+    if arg_ntasks != 1:  # CLI override
+        ntasks = arg_ntasks
+
+    workdir = arg_chdir or directives.get("chdir") or directives.get("D")
+    if not workdir:
+        workdir = os.getcwd()
+
+    hosts = arg_hosts or directives.get("hosts") or ""
+
+    # Resolve executable to absolute path
+    if not os.path.isabs(executable):
+        executable = os.path.abspath(os.path.join(workdir, executable))
+
+    # Submit job via API using query parameters
+    from urllib.parse import urlencode
+
+    params = {
+        "simulation_name": job_name,
+        "flash_executable_path": executable,
+        "parameter_file_path": "",  # FLASH par file - use default
+        "working_directory": workdir,
+        "ranks_per_node": str(ntasks),
+        "min_nodes": str(nodes),
+        "hosts": hosts,
+    }
+
+    query_string = urlencode(params)
+    result = api_request("POST", f"/flash/launch?{query_string}")
+
+    # Print job submission confirmation
+    if isinstance(result, dict):
+        instance_id_val = result.get("instance_id")
+
+        if instance_id_val is not None:
+            job_id = truncate_id(str(instance_id_val))  # pyright: ignore[reportAny]
+            print(f"Submitted batch job {job_id}")
+        else:
+            # Instance created asynchronously - user should check squeue
+            print("Job submitted successfully")
+            print("Use 'exo squeue' to view job ID")
+    else:
+        print("Job submitted successfully")
+        print("Use 'exo squeue' to view job ID")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
--- a/src/exo/cli/scancel.py
+++ b/src/exo/cli/scancel.py
@@ -0,0 +1,95 @@
+"""scancel - Cancel jobs in the Exo queue.
+
+Usage:
+    exo scancel <jobid> [<jobid>...]
+
+Arguments:
+    jobid   Job ID (or prefix) to cancel. Can specify multiple.
+
+Examples:
+    exo scancel abc123          # Cancel job starting with abc123
+    exo scancel abc123 def456   # Cancel multiple jobs
+"""
+
+import argparse
+import sys
+from typing import Any, cast
+
+from exo.cli.common import api_request, truncate_id
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for scancel command."""
+    parser = argparse.ArgumentParser(
+        prog="exo scancel",
+        description="Cancel jobs in the Exo queue",
+    )
+    parser.add_argument(
+        "jobids",
+        nargs="+",
+        help="Job ID(s) to cancel",
+    )
+
+    parsed = parser.parse_args(args)
+    jobids: list[str] = parsed.jobids  # pyright: ignore[reportAny]
+
+    # Fetch current jobs to resolve partial IDs
+    result = api_request("GET", "/flash/instances")
+    if isinstance(result, list):
+        instances = cast(list[dict[str, Any]], result)
+    else:
+        instances = cast(list[dict[str, Any]], result.get("instances", []))
+
+    # Build lookup of full IDs
+    id_map: dict[str, str] = {}
+    for inst in instances:
+        iid = inst.get("instance_id", "")  # pyright: ignore[reportAny]
+        full_id = str(iid) if iid else ""  # pyright: ignore[reportAny]
+        if full_id:
+            # Map both full ID and truncated versions
+            normalized = full_id.replace("-", "").lower()
+            id_map[normalized] = full_id
+            # Also map prefixes
+            for length in range(4, len(normalized) + 1):
+                prefix = normalized[:length]
+                if prefix not in id_map:
+                    id_map[prefix] = full_id
+
+    cancelled = 0
+    errors = 0
+
+    for jobid in jobids:
+        search = jobid.lower().replace("-", "")
+
+        # Find matching full ID
+        full_id = id_map.get(search)
+        if not full_id:
+            # Try prefix match
+            matches = [fid for key, fid in id_map.items() if key.startswith(search)]
+            if len(matches) == 1:
+                full_id = matches[0]
+            elif len(matches) > 1:
+                print(f"Ambiguous job ID: {jobid} matches multiple jobs")
+                errors += 1
+                continue
+            else:
+                print(f"Job not found: {jobid}")
+                errors += 1
+                continue
+
+        # Cancel the job
+        try:
+            api_request("DELETE", f"/flash/{full_id}")
+            print(f"Job {truncate_id(full_id)} cancelled")
+            cancelled += 1
+        except SystemExit:
+            print(f"Failed to cancel job {truncate_id(full_id)}")
+            errors += 1
+
+    if errors > 0 and cancelled == 0:
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
--- a/src/exo/cli/squeue.py
+++ b/src/exo/cli/squeue.py
@@ -0,0 +1,165 @@
+"""squeue - View the Exo job queue.
+
+Usage:
+    exo squeue [options]
+
+Options:
+    -l, --long      Show detailed output
+    -j, --job ID    Show only this job
+
+Output columns:
+    JOBID   - Job identifier (truncated UUID)
+    NAME    - Job name
+    NODES   - Number of nodes
+    STATE   - Job state (PENDING, RUNNING, FAILED, etc.)
+"""
+
+import argparse
+import sys
+from typing import Any, cast
+
+from exo.cli.common import api_request, format_table, truncate_id
+
+# Map Exo runner statuses to SLURM-like states
+STATUS_MAP: dict[str, str] = {
+    "RunnerIdle": "PENDING",
+    "RunnerConnecting": "CONFIGURING",
+    "RunnerConnected": "CONFIGURING",
+    "RunnerLoading": "CONFIGURING",
+    "RunnerLoaded": "CONFIGURING",
+    "RunnerWarmingUp": "CONFIGURING",
+    "RunnerReady": "COMPLETING",
+    "RunnerRunning": "RUNNING",
+    "RunnerShuttingDown": "COMPLETING",
+    "RunnerShutdown": "COMPLETED",
+    "RunnerFailed": "FAILED",
+}
+
+
+def get_job_state(runner_statuses: dict[str, Any]) -> str:
+    """Determine overall job state from runner statuses."""
+    if not runner_statuses:
+        return "PENDING"
+
+    states: set[str] = set()
+    for status_val in runner_statuses.values():  # pyright: ignore[reportAny]
+        if isinstance(status_val, dict):
+            # Extract status type from discriminated union
+            type_val = status_val.get("type", "RunnerIdle")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            status_type = str(type_val) if type_val else "RunnerIdle"  # pyright: ignore[reportUnknownArgumentType]
+        elif isinstance(status_val, str):
+            status_type = status_val
+        else:
+            status_type = "RunnerIdle"
+        # Strip parentheses from status strings like "RunnerRunning()"
+        if status_type.endswith("()"):
+            status_type = status_type[:-2]
+        states.add(STATUS_MAP.get(status_type, "UNKNOWN"))
+
+    # Priority order for overall state
+    if "FAILED" in states:
+        return "FAILED"
+    if "RUNNING" in states:
+        return "RUNNING"
+    if "CONFIGURING" in states:
+        return "CONFIGURING"
+    if "COMPLETING" in states:
+        return "COMPLETING"
+    if "COMPLETED" in states:
+        return "COMPLETED"
+    if "PENDING" in states:
+        return "PENDING"
+    return "UNKNOWN"
+
+
+def main(args: list[str]) -> int:
+    """Main entry point for squeue command."""
+    parser = argparse.ArgumentParser(
+        prog="exo squeue",
+        description="View the Exo job queue",
+    )
+    parser.add_argument(
+        "-l",
+        "--long",
+        action="store_true",
+        help="Show detailed output",
+    )
+    parser.add_argument(
+        "-j",
+        "--job",
+        help="Show only this job ID",
+    )
+
+    parsed = parser.parse_args(args)
+
+    # Extract typed values
+    long_format: bool = parsed.long  # pyright: ignore[reportAny]
+    job_filter: str | None = parsed.job  # pyright: ignore[reportAny]
+
+    # Fetch jobs from API - returns list directly
+    result = api_request("GET", "/flash/instances")
+    # API returns list directly, not {"instances": [...]}
+    if isinstance(result, list):
+        instances = cast(list[dict[str, Any]], result)
+    else:
+        instances = cast(list[dict[str, Any]], result.get("instances", []))
+
+    if not instances:
+        # No jobs - just print header
+        if long_format:
+            print("JOBID           NAME            NODES  RANKS  STATE        WORKDIR")
+        else:
+            print("JOBID       NAME            NODES  STATE")
+        return 0
+
+    # Filter by job ID if specified
+    if job_filter:
+        search = job_filter.lower()
+        filtered: list[dict[str, Any]] = []
+        for i in instances:
+            iid = i.get("instance_id", "")  # pyright: ignore[reportAny]
+            if search in str(iid).lower().replace("-", ""):  # pyright: ignore[reportAny]
+                filtered.append(i)
+        instances = filtered
+
+    # Build table
+    rows: list[list[str]] = []
+
+    if long_format:
+        headers = ["JOBID", "NAME", "NODES", "RANKS", "STATE", "WORKDIR"]
+        for inst in instances:
+            iid_val = inst.get("instance_id", "")  # pyright: ignore[reportAny]
+            instance_id = str(iid_val) if iid_val else ""  # pyright: ignore[reportAny]
+            job_id = truncate_id(instance_id, 12)
+            name_val = inst.get("simulation_name", "")  # pyright: ignore[reportAny]
+            name = (str(name_val) if name_val else "")[:15]  # pyright: ignore[reportAny]
+            runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
+            nodes = str(len(runner_statuses))
+            ranks_val = inst.get("total_ranks", 0)  # pyright: ignore[reportAny]
+            ranks = str(ranks_val) if ranks_val else "0"  # pyright: ignore[reportAny]
+            state = get_job_state(runner_statuses)
+            workdir_val = inst.get("working_directory", "")  # pyright: ignore[reportAny]
+            workdir = str(workdir_val) if workdir_val else ""  # pyright: ignore[reportAny]
+            # Truncate workdir for display
+            if len(workdir) > 30:
+                workdir = "..." + workdir[-27:]
+            rows.append([job_id, name, nodes, ranks, state, workdir])
+    else:
+        headers = ["JOBID", "NAME", "NODES", "STATE"]
+        for inst in instances:
+            iid_val = inst.get("instance_id", "")  # pyright: ignore[reportAny]
+            instance_id = str(iid_val) if iid_val else ""  # pyright: ignore[reportAny]
+            job_id = truncate_id(instance_id, 8)
+            name_val = inst.get("simulation_name", "")  # pyright: ignore[reportAny]
+            name = (str(name_val) if name_val else "")[:15]  # pyright: ignore[reportAny]
+            runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
+            nodes = str(len(runner_statuses))
+            state = get_job_state(runner_statuses)
+            rows.append([job_id, name, nodes, state])
+
+    print(format_table(headers, rows))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -195,6 +195,14 @@ class Node:


 def main():
+    # Check for SLURM-compatible subcommands first
+    import sys
+
+    if len(sys.argv) > 1 and sys.argv[1] in ("sbatch", "squeue", "scancel", "salloc"):
+        from exo.cli import run_subcommand
+
+        sys.exit(run_subcommand(sys.argv[1], sys.argv[2:]))
+
    args = Args.parse()
    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
    resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
@@ -205,6 +213,11 @@ def main():
    logger.info("Starting EXO")
    logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")

+    # Discover and register plugins
+    from exo.plugins.registry import discover_plugins
+
+    discover_plugins()
+
    # Set FAST_SYNCH override env var for runner subprocesses
    if args.fast_synch is True:
        os.environ["EXO_FAST_SYNCH"] = "on"
--- a/src/exo/master/adapters/init.py
+++ b/src/exo/master/adapters/init.py
@@ -1 +0,0 @@
-"""API adapters for different API formats (Claude, OpenAI Responses, etc.)."""
--- a/src/exo/master/adapters/chat_completions.py
+++ b/src/exo/master/adapters/chat_completions.py
@@ -1,186 +0,0 @@
-"""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,
-            )
-        ],
-    )
--- a/src/exo/master/adapters/claude.py
+++ b/src/exo/master/adapters/claude.py
@@ -1,190 +0,0 @@
-"""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"
--- a/src/exo/master/adapters/responses.py
+++ b/src/exo/master/adapters/responses.py
@@ -1,173 +0,0 @@
-"""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"
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -1,7 +1,9 @@
+import asyncio
+import os
 import time
 from collections.abc import AsyncGenerator
 from http import HTTPStatus
-from typing import cast
+from typing import Any, Optional, cast

 import anyio
 from anyio import BrokenResourceError, create_task_group
@@ -14,21 +16,8 @@ from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType
 from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger
+from pydantic import BaseModel

-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
@@ -41,7 +30,6 @@ from exo.shared.types.api import (
    ChatCompletionChoice,
    ChatCompletionMessage,
    ChatCompletionResponse,
-    ChatCompletionTaskParams,
    CreateInstanceParams,
    CreateInstanceResponse,
    DeleteInstanceResponse,
@@ -54,12 +42,9 @@ 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,
@@ -78,12 +63,13 @@ from exo.shared.types.events import (
 )
 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.worker.instances import Instance, InstanceId, InstanceMeta
+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
 from exo.utils.channels import Receiver, Sender, channel
@@ -91,6 +77,39 @@ from exo.utils.dashboard_path import find_dashboard
 from exo.utils.event_buffer import OrderedBuffer


+class ExecuteRequest(BaseModel):
+    """Request to execute a command."""
+
+    command: list[str]
+    cwd: Optional[str] = None
+    env: Optional[dict[str, str]] = None
+
+
+class ExecuteResponse(BaseModel):
+    """Response from command execution."""
+
+    exit_code: int
+    stdout: str
+    stderr: str
+
+
+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_meta(model_id: str) -> ModelMetadata:
    if model_id in MODEL_CARDS:
        model_card = MODEL_CARDS[model_id]
@@ -130,6 +149,7 @@ class API:
        self._setup_exception_handlers()
        self._setup_cors()
        self._setup_routes()
+        self._register_plugin_routes()

        self.app.mount(
            "/",
@@ -196,10 +216,60 @@ 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)
+        # Remote execution endpoint (used by exo-rsh for MPI)
+        self.app.post("/execute")(self.execute)
+
+    def _register_plugin_routes(self) -> None:
+        """Register API routes from all loaded plugins."""
+        import functools
+        import inspect
+
+        from exo.plugins.context import PluginContext
+        from exo.plugins.registry import PluginRegistry
+
+        registry = PluginRegistry.get()
+
+        for method, path, handler, plugin in registry.get_all_api_routes():
+            # Create a wrapper that injects the plugin context while preserving
+            # the original function signature for FastAPI parameter extraction
+            @functools.wraps(handler)
+            async def wrapped_handler(  # pyright: ignore[reportAny]
+                *args: Any,  # pyright: ignore[reportAny]
+                _handler: Any = handler,  # pyright: ignore[reportAny]
+                **kwargs: Any,  # pyright: ignore[reportAny]
+            ) -> Any:
+                context = PluginContext(
+                    state=self.state,
+                    send_command=self._send,
+                    node_id=self.node_id,
+                )
+                # Pass context as first argument, then forward all other args
+                return await _handler(context, *args, **kwargs)  # pyright: ignore[reportAny]
+
+            # Modify the wrapper signature to match the original handler
+            # but without the 'ctx' parameter (we inject it)
+            orig_sig = inspect.signature(handler)
+            params = list(orig_sig.parameters.values())
+            # Remove the first parameter (ctx: PluginContext)
+            if params and params[0].name in ("ctx", "context"):
+                params = params[1:]
+            wrapped_handler.__signature__ = orig_sig.replace(parameters=params)  # pyright: ignore[reportAttributeAccessIssue]
+
+            # Register the route based on HTTP method
+            if method == "get":
+                self.app.get(path)(wrapped_handler)
+            elif method == "post":
+                self.app.post(path)(wrapped_handler)
+            elif method == "delete":
+                self.app.delete(path)(wrapped_handler)
+            elif method == "put":
+                self.app.put(path)(wrapped_handler)
+
+            logger.debug(
+                f"Registered plugin route: {method.upper()} {path} ({plugin.name})"
+            )

    async def place_instance(self, payload: PlaceInstanceParams):
        command = PlaceInstance(
@@ -258,6 +328,8 @@ class API:
                    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,
            )
@@ -313,6 +385,8 @@ class API:
                            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,
                    )
@@ -403,13 +477,11 @@ class API:
            instance_id=instance_id,
        )

-    async def _token_chunk_stream(
+    async def _chat_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]()

@@ -432,6 +504,77 @@ 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:
@@ -441,7 +584,7 @@ class API:

        stats: GenerationStats | None = None

-        async for chunk in self._token_chunk_stream(command_id):
+        async for chunk in self._chat_chunk_stream(command_id):
            if chunk.finish_reason == "error":
                raise HTTPException(
                    status_code=500,
@@ -485,169 +628,60 @@ class API:
    async def chat_completions(
        self, payload: ChatCompletionTaskParams
    ) -> ChatCompletionResponse | StreamingResponse:
-        """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
+        """Handle chat completions, supporting both streaming and non-streaming responses."""
+        model_meta = await resolve_model_meta(payload.model)
+        payload.model = model_meta.model_id

        if not any(
-            instance.shard_assignments.model_id == internal_params.model
+            instance.shard_assignments.model_id == payload.model
            for instance in self.state.instances.values()
        ):
-            await self._trigger_notify_user_to_download_model(internal_params.model)
+            await self._trigger_notify_user_to_download_model(payload.model)
            raise HTTPException(
-                status_code=404,
-                detail=f"No instance found for model {internal_params.model}",
+                status_code=404, detail=f"No instance found for model {payload.model}"
            )

-        command = ChatCompletion(request_params=internal_params)
+        command = ChatCompletion(
+            request_params=payload,
+        )
        await self._send(command)
-
        if payload.stream:
            return StreamingResponse(
-                generate_chat_stream(
-                    command.command_id,
-                    self._token_chunk_stream(command.command_id),
-                ),
+                self._generate_chat_stream(command.command_id),
                media_type="text/event-stream",
-                headers={
-                    "Cache-Control": "no-cache",
-                    "Connection": "close",
-                    "X-Accel-Buffering": "no",
-                },
            )

-        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
+        return await self._collect_chat_completion(command.command_id)

    async def bench_chat_completions(
        self, payload: BenchChatCompletionTaskParams
    ) -> BenchChatCompletionResponse:
-        # 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
+        model_meta = await resolve_model_meta(payload.model)
+        payload.model = model_meta.model_id

        if not any(
-            instance.shard_assignments.model_id == internal_params.model
+            instance.shard_assignments.model_id == payload.model
            for instance in self.state.instances.values()
        ):
-            await self._trigger_notify_user_to_download_model(internal_params.model)
+            await self._trigger_notify_user_to_download_model(payload.model)
            raise HTTPException(
-                status_code=404,
-                detail=f"No instance found for model {internal_params.model}",
+                status_code=404, detail=f"No instance found for model {payload.model}"
            )

-        internal_params.stream = False
+        payload.stream = False

-        command = ChatCompletion(request_params=internal_params)
+        command = ChatCompletion(request_params=payload)
        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 node in self.state.topology.list_nodes():
-            if node.node_profile is not None:
-                total_available += node.node_profile.memory.ram_available
+        for memory in self.state.node_memory.values():
+            total_available += memory.ram_available

        return total_available

@@ -668,6 +702,65 @@ class API:
            ]
        )

+    async def execute(self, request: ExecuteRequest) -> ExecuteResponse:
+        """Execute a command locally. Used by exo-rsh for MPI remote execution."""
+        cmd_str = " ".join(request.command)
+        logger.info(f"Executing: {cmd_str}")
+
+        try:
+            # Build environment
+            env = os.environ.copy()
+            if request.env:
+                env.update(request.env)
+
+            # Check if command contains shell metacharacters
+            # If so, run through shell. mpirun sends complex commands like:
+            # "VAR=value;export VAR;/path/to/prted --args"
+            needs_shell = any(c in cmd_str for c in ";|&$`")
+
+            if needs_shell:
+                process = await asyncio.create_subprocess_shell(
+                    cmd_str,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    cwd=request.cwd,
+                    env=env,
+                )
+            else:
+                process = await asyncio.create_subprocess_exec(
+                    *request.command,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    cwd=request.cwd,
+                    env=env,
+                )
+
+            stdout, stderr = await process.communicate()
+            exit_code = process.returncode or 0
+
+            logger.info(f"Command completed with exit code {exit_code}")
+
+            return ExecuteResponse(
+                exit_code=exit_code,
+                stdout=stdout.decode("utf-8", errors="replace"),
+                stderr=stderr.decode("utf-8", errors="replace"),
+            )
+
+        except FileNotFoundError:
+            logger.error(f"Command not found: {request.command[0]}")
+            return ExecuteResponse(
+                exit_code=127,
+                stdout="",
+                stderr=f"Command not found: {request.command[0]}",
+            )
+        except Exception as e:
+            logger.error(f"Execution error: {e}")
+            return ExecuteResponse(
+                exit_code=1,
+                stdout="",
+                stderr=str(e),
+            )
+
    async def run(self):
        cfg = Config()
        cfg.bind = f"0.0.0.0:{self.port}"
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -27,6 +27,7 @@ from exo.shared.types.events import (
    ForwarderEvent,
    IndexedEvent,
    InstanceDeleted,
+    NodeGatheredInfo,
    NodeTimedOut,
    TaskCreated,
    TaskDeleted,
@@ -95,102 +96,128 @@ class Master:
        self._tg.cancel_scope.cancel()

    async def _command_processor(self) -> None:
+        from exo.plugins.registry import PluginRegistry
+
+        registry = PluginRegistry.get()
+
        with self.command_receiver as commands:
            async for forwarder_command in commands:
                try:
                    logger.info(f"Executing command: {forwarder_command.command}")
                    generated_events: list[Event] = []
                    command = forwarder_command.command
-                    match command:
-                        case TestCommand():
-                            pass
-                        case ChatCompletion():
-                            instance_task_counts: dict[InstanceId, int] = {}
-                            for instance in self.state.instances.values():
-                                if (
-                                    instance.shard_assignments.model_id
-                                    == command.request_params.model
-                                ):
-                                    task_count = sum(
-                                        1
-                                        for task in self.state.tasks.values()
-                                        if task.instance_id == instance.instance_id
-                                    )
-                                    instance_task_counts[instance.instance_id] = (
-                                        task_count
+
+                    # Check if a plugin handles this command
+                    plugin = registry.get_plugin_for_command(command)
+                    if plugin is not None:
+                        events = plugin.process_command(
+                            command,
+                            self.state.topology,
+                            self.state.instances,
+                        )
+                        generated_events.extend(events)
+                    else:
+                        # Core command handling
+                        match command:
+                            case TestCommand():
+                                pass
+                            case ChatCompletion():
+                                instance_task_counts: dict[InstanceId, int] = {}
+                                for instance in self.state.instances.values():
+                                    if (
+                                        instance.shard_assignments.model_id
+                                        == command.request_params.model
+                                    ):
+                                        task_count = sum(
+                                            1
+                                            for task in self.state.tasks.values()
+                                            if task.instance_id == instance.instance_id
+                                        )
+                                        instance_task_counts[instance.instance_id] = (
+                                            task_count
+                                        )
+
+                                if not instance_task_counts:
+                                    raise ValueError(
+                                        f"No instance found for model {command.request_params.model}"
                                    )

-                            if not instance_task_counts:
-                                raise ValueError(
-                                    f"No instance found for model {command.request_params.model}"
+                                available_instance_ids = sorted(
+                                    instance_task_counts.keys(),
+                                    key=lambda instance_id: instance_task_counts[
+                                        instance_id
+                                    ],
                                )

-                            available_instance_ids = sorted(
-                                instance_task_counts.keys(),
-                                key=lambda instance_id: instance_task_counts[
-                                    instance_id
-                                ],
-                            )
-
-                            task_id = TaskId()
-                            generated_events.append(
-                                TaskCreated(
-                                    task_id=task_id,
-                                    task=ChatCompletionTask(
+                                task_id = TaskId()
+                                generated_events.append(
+                                    TaskCreated(
                                        task_id=task_id,
-                                        command_id=command.command_id,
-                                        instance_id=available_instance_ids[0],
-                                        task_status=TaskStatus.Pending,
-                                        task_params=command.request_params,
-                                    ),
+                                        task=ChatCompletionTask(
+                                            task_id=task_id,
+                                            command_id=command.command_id,
+                                            instance_id=available_instance_ids[0],
+                                            task_status=TaskStatus.Pending,
+                                            task_params=command.request_params,
+                                        ),
+                                    )
                                )
-                            )

-                            self.command_task_mapping[command.command_id] = task_id
-                        case DeleteInstance():
-                            placement = delete_instance(command, self.state.instances)
-                            transition_events = get_transition_events(
-                                self.state.instances, placement
-                            )
-                            generated_events.extend(transition_events)
-                        case PlaceInstance():
-                            placement = place_instance(
-                                command,
-                                self.state.topology,
-                                self.state.instances,
-                            )
-                            transition_events = get_transition_events(
-                                self.state.instances, placement
-                            )
-                            generated_events.extend(transition_events)
-                        case CreateInstance():
-                            placement = add_instance_to_placements(
-                                command,
-                                self.state.topology,
-                                self.state.instances,
-                            )
-                            transition_events = get_transition_events(
-                                self.state.instances, placement
-                            )
-                            generated_events.extend(transition_events)
-                        case TaskFinished():
-                            generated_events.append(
-                                TaskDeleted(
-                                    task_id=self.command_task_mapping[
+                                self.command_task_mapping[command.command_id] = task_id
+                            case DeleteInstance():
+                                placement = delete_instance(
+                                    command, self.state.instances
+                                )
+                                transition_events = get_transition_events(
+                                    self.state.instances, placement
+                                )
+                                generated_events.extend(transition_events)
+                            case PlaceInstance():
+                                placement = place_instance(
+                                    command,
+                                    self.state.topology,
+                                    self.state.instances,
+                                    self.state.node_memory,
+                                    self.state.node_network,
+                                )
+                                transition_events = get_transition_events(
+                                    self.state.instances, placement
+                                )
+                                generated_events.extend(transition_events)
+                            case CreateInstance():
+                                placement = add_instance_to_placements(
+                                    command,
+                                    self.state.topology,
+                                    self.state.instances,
+                                )
+                                transition_events = get_transition_events(
+                                    self.state.instances, placement
+                                )
+                                generated_events.extend(transition_events)
+                            case TaskFinished():
+                                generated_events.append(
+                                    TaskDeleted(
+                                        task_id=self.command_task_mapping[
+                                            command.finished_command_id
+                                        ]
+                                    )
+                                )
+                                if (
+                                    command.finished_command_id
+                                    in self.command_task_mapping
+                                ):
+                                    del self.command_task_mapping[
                                        command.finished_command_id
                                    ]
-                                )
-                            )
-                            if command.finished_command_id in self.command_task_mapping:
-                                del self.command_task_mapping[
-                                    command.finished_command_id
-                                ]
-                        case RequestEventLog():
-                            # We should just be able to send everything, since other buffers will ignore old messages
-                            for i in range(command.since_idx, len(self._event_log)):
-                                await self._send_event(
-                                    IndexedEvent(idx=i, event=self._event_log[i])
-                                )
+                            case RequestEventLog():
+                                # We should just be able to send everything, since other buffers will ignore old messages
+                                for i in range(command.since_idx, len(self._event_log)):
+                                    await self._send_event(
+                                        IndexedEvent(idx=i, event=self._event_log[i])
+                                    )
+                            case _:
+                                # Plugin-managed commands are handled above
+                                pass
                    for event in generated_events:
                        await self.event_sender.send(event)
                except ValueError as e:
@@ -200,9 +227,7 @@ class Master:
    async def _plan(self) -> None:
        while True:
            # kill broken instances
-            connected_node_ids = set(
-                [x.node_id for x in self.state.topology.list_nodes()]
-            )
+            connected_node_ids = set(self.state.topology.list_nodes())
            for instance_id, instance in self.state.instances.items():
                for node_id in instance.shard_assignments.node_to_runner:
                    if node_id not in connected_node_ids:
@@ -237,6 +262,8 @@ class Master:
                    self.state = apply(self.state, indexed)

                    event._master_time_stamp = datetime.now(tz=timezone.utc)  # pyright: ignore[reportPrivateUsage]
+                    if isinstance(event, NodeGatheredInfo):
+                        event.when = str(datetime.now(tz=timezone.utc))

                    self._event_log.append(event)
                    await self._send_event(indexed)
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -6,9 +6,10 @@ 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,
@@ -19,10 +20,11 @@ from exo.shared.types.commands import (
    DeleteInstance,
    PlaceInstance,
 )
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId
-from exo.shared.types.topology import NodeInfo
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.worker.instances import (
    Instance,
    InstanceId,
@@ -52,19 +54,15 @@ def place_instance(
    command: PlaceInstance,
    topology: Topology,
    current_instances: Mapping[InstanceId, Instance],
+    node_memory: Mapping[NodeId, MemoryUsage],
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[InstanceId, Instance]:
-    all_nodes = list(topology.list_nodes())
-
-    logger.info("finding cycles:")
    cycles = topology.get_cycles()
-    singleton_cycles = [[node] for node in all_nodes]
-    candidate_cycles = list(
-        filter(lambda it: len(it) >= command.min_nodes, cycles + singleton_cycles)
-    )
+    candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
    cycles_with_sufficient_memory = filter_cycles_by_memory(
-        candidate_cycles, command.model_meta.storage_size
+        candidate_cycles, node_memory, command.model_meta.storage_size
    )
-    if not cycles_with_sufficient_memory:
+    if len(cycles_with_sufficient_memory) == 0:
        raise ValueError("No cycles found with sufficient memory")

    if command.sharding == Sharding.Tensor:
@@ -92,44 +90,38 @@ def place_instance(
    smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)

    smallest_tb_cycles = [
-        cycle
-        for cycle in smallest_cycles
-        if topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
+        cycle for cycle in smallest_cycles if topology.is_thunderbolt_cycle(cycle)
    ]

    if smallest_tb_cycles != []:
        smallest_cycles = smallest_tb_cycles

-    cycles_with_leaf_nodes: list[list[NodeInfo]] = [
+    cycles_with_leaf_nodes: list[Cycle] = [
        cycle
        for cycle in smallest_cycles
-        if any(topology.node_is_leaf(node.node_id) for node in cycle)
+        if any(topology.node_is_leaf(node_id) for node_id in cycle)
    ]

    selected_cycle = max(
        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
        key=lambda cycle: sum(
-            (
-                node.node_profile.memory.ram_available
-                for node in cycle
-                if node.node_profile is not None
-            ),
+            (node_memory[node_id].ram_available for node_id in cycle),
            start=Memory(),
        ),
    )

    shard_assignments = get_shard_assignments(
-        command.model_meta, selected_cycle, command.sharding
+        command.model_meta, selected_cycle, command.sharding, node_memory
    )

-    cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle)
+    cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)

    instance_id = InstanceId()
    target_instances = dict(deepcopy(current_instances))

    if len(selected_cycle) == 1:
        logger.warning(
-            "You have likely selected ibv for a single node instance; falling back to MlxRing"
+            "You have likely selected jaccl for a single node instance; falling back to MlxRing"
        )

        command.instance_meta = InstanceMeta.MlxRing
@@ -137,19 +129,20 @@ def place_instance(
    # TODO: Single node instances
    match command.instance_meta:
        case InstanceMeta.MlxJaccl:
-            mlx_ibv_devices = get_mlx_ibv_devices_matrix(
-                selected_cycle,
+            mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
+                [node_id for node_id in selected_cycle],
                cycle_digraph,
            )
            mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
-                selected_cycle,
+                coordinator=selected_cycle.node_ids[0],
                coordinator_port=random_ephemeral_port(),
                cycle_digraph=cycle_digraph,
+                node_network=node_network,
            )
            target_instances[instance_id] = MlxJacclInstance(
                instance_id=instance_id,
                shard_assignments=shard_assignments,
-                ibv_devices=mlx_ibv_devices,
+                jaccl_devices=mlx_jaccl_devices,
                jaccl_coordinators=mlx_jaccl_coordinators,
            )
        case InstanceMeta.MlxRing:
@@ -158,6 +151,7 @@ def place_instance(
                selected_cycle=selected_cycle,
                cycle_digraph=cycle_digraph,
                ephemeral_port=ephemeral_port,
+                node_network=node_network,
            )
            target_instances[instance_id] = MlxRingInstance(
                instance_id=instance_id,
@@ -165,6 +159,11 @@ def place_instance(
                hosts_by_node=hosts_by_node,
                ephemeral_port=ephemeral_port,
            )
+        case _:
+            # Plugin-managed instance types have their own placement functions
+            raise ValueError(
+                f"Instance type {command.instance_meta} must use plugin placement"
+            )

    return target_instances

--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -1,15 +1,13 @@
-from collections.abc import Generator
-from typing import TypeGuard, cast
+from collections.abc import Generator, Mapping

 from loguru import logger
-from pydantic import BaseModel

 from exo.shared.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelMetadata
-from exo.shared.types.profiling import NodePerformanceProfile
-from exo.shared.types.topology import NodeInfo
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
    PipelineShardMetadata,
@@ -19,32 +17,28 @@ from exo.shared.types.worker.shards import (
 )


-class NodeWithProfile(BaseModel):
-    node_id: NodeId
-    node_profile: NodePerformanceProfile
-
-
-def narrow_all_nodes(nodes: list[NodeInfo]) -> TypeGuard[list[NodeWithProfile]]:
-    return all(node.node_profile is not None for node in nodes)
-
-
 def filter_cycles_by_memory(
-    cycles: list[list[NodeInfo]], required_memory: Memory
-) -> list[list[NodeInfo]]:
-    filtered_cycles: list[list[NodeInfo]] = []
+    cycles: list[Cycle],
+    node_memory: Mapping[NodeId, MemoryUsage],
+    required_memory: Memory,
+) -> list[Cycle]:
+    filtered_cycles: list[Cycle] = []
    for cycle in cycles:
-        if not narrow_all_nodes(cycle):
+        if not all(node in node_memory for node in cycle):
            continue

        total_mem = sum(
-            (node.node_profile.memory.ram_available for node in cycle), start=Memory()
+            (node_memory[node_id].ram_available for node_id in cycle.node_ids),
+            start=Memory(),
        )
        if total_mem >= required_memory:
-            filtered_cycles.append(cast(list[NodeInfo], cycle))
+            filtered_cycles.append(cycle)
    return filtered_cycles


-def get_smallest_cycles(cycles: list[list[NodeInfo]]) -> list[list[NodeInfo]]:
+def get_smallest_cycles(
+    cycles: list[Cycle],
+) -> list[Cycle]:
    min_nodes = min(len(cycle) for cycle in cycles)
    return [cycle for cycle in cycles if len(cycle) == min_nodes]

@@ -82,13 +76,14 @@ def allocate_layers_proportionally(

 def get_shard_assignments_for_pipeline_parallel(
    model_meta: ModelMetadata,
-    selected_cycle: list[NodeWithProfile],
+    cycle: Cycle,
+    node_memory: Mapping[NodeId, MemoryUsage],
 ):
-    if not selected_cycle:
+    if not cycle.node_ids:
        raise ValueError("Cannot create shard assignments for empty node cycle")

    cycle_memory = sum(
-        (node.node_profile.memory.ram_available for node in selected_cycle),
+        (node_memory[node_id].ram_available for node_id in cycle.node_ids),
        start=Memory(),
    )

@@ -96,35 +91,35 @@ def get_shard_assignments_for_pipeline_parallel(
        raise ValueError("Cannot create shard assignments: total available memory is 0")

    total_layers = model_meta.n_layers
-    world_size = len(selected_cycle)
+    world_size = len(cycle)
    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
    node_to_runner: dict[NodeId, RunnerId] = {}

    layer_allocations = allocate_layers_proportionally(
        total_layers=total_layers,
        memory_fractions=[
-            node.node_profile.memory.ram_available.in_bytes / cycle_memory.in_bytes
-            for node in selected_cycle
+            node_memory[node_id].ram_available.in_bytes / cycle_memory.in_bytes
+            for node_id in cycle.node_ids
        ],
    )

    # Validate each node has sufficient memory for its assigned layers
    memory_per_layer = model_meta.storage_size.in_bytes / total_layers
-    for i, (node, node_layers) in enumerate(
-        zip(selected_cycle, layer_allocations, strict=True)
+    for i, (node_id, node_layers) in enumerate(
+        zip(cycle.node_ids, layer_allocations, strict=True)
    ):
        required_memory = node_layers * memory_per_layer
-        available_memory = node.node_profile.memory.ram_available.in_bytes
+        available_memory = node_memory[node_id].ram_available.in_bytes
        if required_memory > available_memory:
            raise ValueError(
-                f"Node {i} ({node.node_id}) has insufficient memory: "
+                f"Node {i} ({node_id}) has insufficient memory: "
                f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
                f"but only has {available_memory / (1024**3):.2f} GB available"
            )

    layers_assigned = 0
-    for i, (node, node_layers) in enumerate(
-        zip(selected_cycle, layer_allocations, strict=True)
+    for i, (node_id, node_layers) in enumerate(
+        zip(cycle.node_ids, layer_allocations, strict=True)
    ):
        runner_id = RunnerId()

@@ -138,7 +133,7 @@ def get_shard_assignments_for_pipeline_parallel(
        )

        runner_to_shard[runner_id] = shard
-        node_to_runner[node.node_id] = runner_id
+        node_to_runner[node_id] = runner_id
        layers_assigned += node_layers

    shard_assignments = ShardAssignments(
@@ -152,14 +147,14 @@ def get_shard_assignments_for_pipeline_parallel(

 def get_shard_assignments_for_tensor_parallel(
    model_meta: ModelMetadata,
-    selected_cycle: list[NodeWithProfile],
+    cycle: Cycle,
 ):
    total_layers = model_meta.n_layers
-    world_size = len(selected_cycle)
+    world_size = len(cycle)
    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
    node_to_runner: dict[NodeId, RunnerId] = {}

-    for i, node in enumerate(selected_cycle):
+    for i, node_id in enumerate(cycle):
        shard = TensorShardMetadata(
            model_meta=model_meta,
            device_rank=i,
@@ -172,7 +167,7 @@ def get_shard_assignments_for_tensor_parallel(
        runner_id = RunnerId()

        runner_to_shard[runner_id] = shard
-        node_to_runner[node.node_id] = runner_id
+        node_to_runner[node_id] = runner_id

    shard_assignments = ShardAssignments(
        model_id=model_meta.model_id,
@@ -185,21 +180,21 @@ def get_shard_assignments_for_tensor_parallel(

 def get_shard_assignments(
    model_meta: ModelMetadata,
-    selected_cycle: list[NodeInfo],
+    cycle: Cycle,
    sharding: Sharding,
+    node_memory: Mapping[NodeId, MemoryUsage],
 ) -> ShardAssignments:
-    if not narrow_all_nodes(selected_cycle):
-        raise ValueError("All nodes must have profiles to create shard assignments")
    match sharding:
        case Sharding.Pipeline:
            return get_shard_assignments_for_pipeline_parallel(
                model_meta=model_meta,
-                selected_cycle=selected_cycle,
+                cycle=cycle,
+                node_memory=node_memory,
            )
        case Sharding.Tensor:
            return get_shard_assignments_for_tensor_parallel(
                model_meta=model_meta,
-                selected_cycle=selected_cycle,
+                cycle=cycle,
            )


@@ -214,38 +209,40 @@ def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
            )
        return []

+    cycle = cycles[0]
+
    get_thunderbolt = False
-    if cycle_digraph.is_thunderbolt_cycle(cycles[0]):
+    if cycle_digraph.is_thunderbolt_cycle(cycle):
        get_thunderbolt = True

    logger.info(f"Using thunderbolt cycle: {get_thunderbolt}")

-    cycle = cycles[0]
    hosts: list[Host] = []
    for i in range(len(cycle)):
-        current_node = cycle[i]
-        next_node = cycle[(i + 1) % len(cycle)]
+        current_node = cycle.node_ids[i]
+        next_node = cycle.node_ids[(i + 1) % len(cycle)]

-        for connection in cycle_digraph.list_connections():
-            if (
-                connection.local_node_id == current_node.node_id
-                and connection.send_back_node_id == next_node.node_id
-            ):
-                if get_thunderbolt and not connection.is_thunderbolt():
-                    continue
-                assert connection.send_back_multiaddr is not None
-                host = Host(
-                    ip=connection.send_back_multiaddr.ip_address,
-                    port=connection.send_back_multiaddr.port,
-                )
-                hosts.append(host)
-                break
+        for connection in cycle_digraph.get_all_connections_between(
+            source=current_node, sink=next_node
+        ):
+            if not isinstance(connection, SocketConnection):
+                continue
+
+            if get_thunderbolt and not connection.is_thunderbolt():
+                continue
+
+            host = Host(
+                ip=connection.sink_multiaddr.ip_address,
+                port=connection.sink_multiaddr.port,
+            )
+            hosts.append(host)
+            break

    return hosts


-def get_mlx_ibv_devices_matrix(
-    selected_cycle: list[NodeInfo],
+def get_mlx_jaccl_devices_matrix(
+    selected_cycle: list[NodeId],
    cycle_digraph: Topology,
 ) -> list[list[str | None]]:
    """Build connectivity matrix mapping device i to device j via RDMA interface names.
@@ -264,72 +261,37 @@ def get_mlx_ibv_devices_matrix(
            if i == j:
                continue

-            # Find the IP J uses to talk to I
-            for connection_ip, _ in _find_connection_ip(node_j, node_i, cycle_digraph):
-                # This is a local IP on I, which is attached to an interface: find that interface
-                if interface_name := _find_rdma_interface_name_for_ip(
-                    connection_ip, node_i
-                ):
-                    matrix[i][j] = interface_name
-                    logger.info(
-                        f"Interface name for {connection_ip} on {node_i.node_id}: {interface_name}"
-                    )
+            for conn in cycle_digraph.get_all_connections_between(node_i, node_j):
+                if isinstance(conn, RDMAConnection):
+                    matrix[i][j] = conn.source_rdma_iface
                    break
            else:
                logger.warning(
-                    f"Failed to find interface name between {node_i.node_id} and {node_j.node_id}"
+                    f"Failed to find interface name between {node_i} and {node_j}"
                )
                raise ValueError(
-                    "Current ibv backend requires all-to-all rdma connections"
+                    "Current jaccl backend requires all-to-all RDMA connections"
                )

    return matrix


 def _find_connection_ip(
-    node_i: NodeInfo,
-    node_j: NodeInfo,
+    node_i: NodeId,
+    node_j: NodeId,
    cycle_digraph: Topology,
 ) -> Generator[tuple[str, bool]]:
-    """Find all IP addresses that connect node i to node j, with thunderbolt flag."""
-    for connection in cycle_digraph.list_connections():
-        if (
-            connection.local_node_id == node_i.node_id
-            and connection.send_back_node_id == node_j.node_id
-        ):
-            yield connection.send_back_multiaddr.ip_address, connection.is_thunderbolt()
-
-
-def _find_rdma_interface_name_for_ip(
-    ip_address: str,
-    node_info: NodeInfo,
-) -> str | None:
-    if node_info.node_profile is None:
-        return None
-
-    logger.info(f"Searching {node_info.node_id} for ip {ip_address}:")
-    for interface in node_info.node_profile.network_interfaces:
-        if interface.name not in ["en2", "en3", "en4", "en5", "en6", "en7"]:
-            continue
-        logger.info(f" | {interface.name}: {interface.ip_address}")
-        if interface.ip_address != ip_address:
-            continue
-
-        logger.info("Found")
-        return f"rdma_{interface.name}"
-
-    return None
+    """Find all IP addresses that connect node i to node j."""
+    for connection in cycle_digraph.get_all_connections_between(node_i, node_j):
+        if isinstance(connection, SocketConnection):
+            yield connection.sink_multiaddr.ip_address, connection.is_thunderbolt()


 def _find_interface_name_for_ip(
-    ip_address: str,
-    node_info: NodeInfo,
+    ip_address: str, node_network: NodeNetworkInfo
 ) -> str | None:
    """Find the interface name for an IP address on a node (any interface)."""
-    if node_info.node_profile is None:
-        return None
-
-    for interface in node_info.node_profile.network_interfaces:
+    for interface in node_network.interfaces:
        if interface.ip_address == ip_address:
            return interface.name

@@ -337,7 +299,10 @@ def _find_interface_name_for_ip(


 def _find_ip_prioritised(
-    node: NodeInfo, other_node: NodeInfo, cycle_digraph: Topology
+    node_id: NodeId,
+    other_node_id: NodeId,
+    cycle_digraph: Topology,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> str | None:
    # TODO: Actually prioritize in the correct Ethernet > Wifi > Non-TB > TB order.
    """Find an IP address between nodes with prioritization.
@@ -348,9 +313,14 @@ def _find_ip_prioritised(
    3. Non-Thunderbolt connections
    4. Any other IP address
    """
-    ips = list(_find_connection_ip(node, other_node, cycle_digraph))
+    ips = list(_find_connection_ip(node_id, other_node_id, cycle_digraph))
    # We expect a unique iface -> ip mapping
-    iface_map = {_find_interface_name_for_ip(ip, other_node): ip for ip, _ in ips}
+    iface_map = {
+        _find_interface_name_for_ip(
+            ip, node_network.get(other_node_id, NodeNetworkInfo())
+        ): ip
+        for ip, _ in ips
+    }

    en0_ip = iface_map.get("en0")
    if en0_ip:
@@ -374,9 +344,10 @@ def _find_ip_prioritised(


 def get_mlx_ring_hosts_by_node(
-    selected_cycle: list[NodeInfo],
+    selected_cycle: Cycle,
    cycle_digraph: Topology,
    ephemeral_port: int,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[NodeId, list[Host]]:
    """Generate per-node host lists for MLX ring backend.

@@ -391,14 +362,13 @@ def get_mlx_ring_hosts_by_node(

    hosts_by_node: dict[NodeId, list[Host]] = {}

-    for rank, node in enumerate(selected_cycle):
-        node_id = node.node_id
+    for rank, node_id in enumerate(selected_cycle):
        left_rank = (rank - 1) % world_size
        right_rank = (rank + 1) % world_size

        hosts_for_node: list[Host] = []

-        for idx, other_node in enumerate(selected_cycle):
+        for idx, other_node_id in enumerate(selected_cycle):
            if idx == rank:
                hosts_for_node.append(Host(ip="0.0.0.0", port=ephemeral_port))
                continue
@@ -408,10 +378,12 @@ def get_mlx_ring_hosts_by_node(
                hosts_for_node.append(Host(ip="198.51.100.1", port=0))
                continue

-            connection_ip = _find_ip_prioritised(node, other_node, cycle_digraph)
+            connection_ip = _find_ip_prioritised(
+                node_id, other_node_id, cycle_digraph, node_network
+            )
            if connection_ip is None:
                logger.warning(
-                    f"Failed to find prioritised connection IP between {node_id} and {other_node.node_id}"
+                    f"Failed to find prioritised connection IP between {node_id} and {other_node_id}"
                )
                raise ValueError(
                    "MLX ring backend requires connectivity between neighbouring nodes"
@@ -425,31 +397,34 @@ def get_mlx_ring_hosts_by_node(


 def get_mlx_jaccl_coordinators(
-    selected_cycle: list[NodeInfo],
+    coordinator: NodeId,
    coordinator_port: int,
    cycle_digraph: Topology,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[NodeId, str]:
-    """Get the coordinator addresses for MLX Jaccl (rank 0 device).
+    """Get the coordinator addresses for MLX JACCL (rank 0 device).

    Select an IP address that each node can reach for the rank 0 node. Returns
    address in format "X.X.X.X:PORT" per node.
    """
-    rank_0_node = selected_cycle[0]
-    logger.debug(f"Selecting coordinator from rank 0 node: {rank_0_node.node_id}")
+    logger.info(f"Selecting coordinator: {coordinator}")

-    def get_ip_for_node(n: NodeInfo) -> str:
-        if n.node_id == rank_0_node.node_id:
+    def get_ip_for_node(n: NodeId) -> str:
+        if n == coordinator:
            return "0.0.0.0"

-        ip = _find_ip_prioritised(n, rank_0_node, cycle_digraph)
-        if ip:
+        ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network)
+        if ip is not None:
            return ip

        logger.warning(
-            f"Failed to find directly connected ip between {n.node_id} and {rank_0_node.node_id}"
+            f"Failed to find directly connected ip between {n} and {coordinator}"
+        )
+        raise ValueError(
+            "Current jaccl backend requires all participating devices to be able to communicate"
        )
-        raise ValueError("Current ibv backend requires all-to-all rdma connections")

    return {
-        n.node_id: f"{get_ip_for_node(n)}:{coordinator_port}" for n in selected_cycle
+        n: f"{get_ip_for_node(n)}:{coordinator_port}"
+        for n in cycle_digraph.list_nodes()
    }
--- a/src/exo/master/tests/conftest.py
+++ b/src/exo/master/tests/conftest.py
@@ -1,67 +1,37 @@
-from typing import Callable
-
-import pytest
-
-from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
+    MemoryUsage,
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
 )
-from exo.shared.types.topology import Connection, ConnectionProfile, NodeInfo
+from exo.shared.types.topology import RDMAConnection, SocketConnection


-@pytest.fixture
-def create_node():
-    def _create_node(memory: int, node_id: NodeId | None = None) -> NodeInfo:
-        if node_id is None:
-            node_id = NodeId()
-        return NodeInfo(
-            node_id=node_id,
-            node_profile=NodePerformanceProfile(
-                model_id="test",
-                chip_id="test",
-                friendly_name="test",
-                memory=MemoryPerformanceProfile.from_bytes(
-                    ram_total=1000,
-                    ram_available=memory,
-                    swap_total=1000,
-                    swap_available=1000,
-                ),
-                network_interfaces=[],
-                system=SystemPerformanceProfile(),
-            ),
-        )
-
-    return _create_node
+def create_node_memory(memory: int) -> MemoryUsage:
+    return MemoryUsage.from_bytes(
+        ram_total=1000,
+        ram_available=memory,
+        swap_total=1000,
+        swap_available=1000,
+    )


-# TODO: this is a hack to get the port for the send_back_multiaddr
-@pytest.fixture
-def create_connection() -> Callable[[NodeId, NodeId, int | None], Connection]:
-    port_counter = 1235
-    ip_counter = 1
+def create_node_network() -> NodeNetworkInfo:
+    return NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address=f"169.254.0.{i}")
+            for i in range(10)
+        ]
+    )

-    def _create_connection(
-        source_node_id: NodeId, sink_node_id: NodeId, send_back_port: int | None = None
-    ) -> Connection:
-        nonlocal port_counter
-        nonlocal ip_counter
-        # assign unique ips
-        ip_counter += 1
-        if send_back_port is None:
-            send_back_port = port_counter
-            port_counter += 1
-        return Connection(
-            local_node_id=source_node_id,
-            send_back_node_id=sink_node_id,
-            send_back_multiaddr=Multiaddr(
-                address=f"/ip4/169.254.0.{ip_counter}/tcp/{send_back_port}"
-            ),
-            connection_profile=ConnectionProfile(
-                throughput=1000, latency=1000, jitter=1000
-            ),
-        )

-    return _create_connection
+def create_socket_connection(ip: int, sink_port: int = 1234) -> SocketConnection:
+    return SocketConnection(
+        sink_multiaddr=Multiaddr(address=f"/ip4/169.254.0.{ip}/tcp/{sink_port}"),
+    )
+
+
+def create_rdma_connection(iface: int) -> RDMAConnection:
+    return RDMAConnection(
+        source_rdma_iface=f"rdma_en{iface}", sink_rdma_iface=f"rdma_en{iface}"
+    )
--- a/src/exo/master/tests/test_claude_api.py
+++ b/src/exo/master/tests/test_claude_api.py
@@ -1,283 +0,0 @@
-"""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")
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -7,6 +7,7 @@ from loguru import logger

 from exo.master.main import Master
 from exo.routing.router import get_node_id_keypair
+from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
 from exo.shared.types.commands import (
    ChatCompletion,
    CommandId,
@@ -18,16 +19,13 @@ from exo.shared.types.events import (
    ForwarderEvent,
    IndexedEvent,
    InstanceCreated,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
    TaskCreated,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
+    MemoryUsage,
 )
 from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
 from exo.shared.types.tasks import TaskStatus
@@ -75,29 +73,22 @@ async def test_master():
        tg.start_soon(master.run)

        sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
-        # inject a NodePerformanceProfile event
-        logger.info("inject a NodePerformanceProfile event")
+        # inject a NodeGatheredInfo event
+        logger.info("inject a NodeGatheredInfo event")
        await local_event_sender.send(
            ForwarderEvent(
                origin_idx=0,
                origin=sender_node_id,
                session=session_id,
                event=(
-                    NodePerformanceMeasured(
+                    NodeGatheredInfo(
                        when=str(datetime.now(tz=timezone.utc)),
                        node_id=node_id,
-                        node_profile=NodePerformanceProfile(
-                            model_id="maccy",
-                            chip_id="arm",
-                            friendly_name="test",
-                            memory=MemoryPerformanceProfile(
-                                ram_total=Memory.from_bytes(678948 * 1024),
-                                ram_available=Memory.from_bytes(678948 * 1024),
-                                swap_total=Memory.from_bytes(0),
-                                swap_available=Memory.from_bytes(0),
-                            ),
-                            network_interfaces=[],
-                            system=SystemPerformanceProfile(),
+                        info=MemoryUsage(
+                            ram_total=Memory.from_bytes(678948 * 1024),
+                            ram_available=Memory.from_bytes(678948 * 1024),
+                            swap_total=Memory.from_bytes(0),
+                            swap_available=Memory.from_bytes(0),
                        ),
                    )
                ),
@@ -108,7 +99,7 @@ async def test_master():
        logger.info("wait for initial topology event")
        while len(list(master.state.topology.list_nodes())) == 0:
            await anyio.sleep(0.001)
-        while len(master.state.node_profiles) == 0:
+        while len(master.state.node_memory) == 0:
            await anyio.sleep(0.001)

        logger.info("inject a CreateInstance Command")
@@ -143,9 +134,13 @@ async def test_master():
                command=(
                    ChatCompletion(
                        command_id=CommandId(),
-                        request_params=ResponsesRequest(
+                        request_params=ChatCompletionTaskParams(
                            model="llama-3.2-1b",
-                            input="Hello, how are you?",
+                            messages=[
+                                ChatCompletionMessage(
+                                    role="user", content="Hello, how are you?"
+                                )
+                            ],
                        ),
                    )
                ),
@@ -159,7 +154,7 @@ async def test_master():
        assert events[0].idx == 0
        assert events[1].idx == 1
        assert events[2].idx == 2
-        assert isinstance(events[0].event, NodePerformanceMeasured)
+        assert isinstance(events[0].event, NodeGatheredInfo)
        assert isinstance(events[1].event, InstanceCreated)
        created_instance = events[1].event.instance
        assert isinstance(created_instance, MlxRingInstance)
@@ -196,9 +191,11 @@ 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 == ResponsesRequest(
+        assert events[2].event.task.task_params == ChatCompletionTaskParams(
            model="llama-3.2-1b",
-            input="Hello, how are you?",
+            messages=[
+                ChatCompletionMessage(role="user", content="Hello, how are you?")
+            ],
        )

        await master.shutdown()
--- a/src/exo/master/tests/test_openai_responses_api.py
+++ b/src/exo/master/tests/test_openai_responses_api.py
@@ -1,293 +0,0 @@
-"""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")
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -1,20 +1,24 @@
-from typing import Callable
-
 import pytest
-from loguru import logger

 from exo.master.placement import (
    get_transition_events,
    place_instance,
 )
+from exo.master.tests.conftest import (
+    create_node_memory,
+    create_node_network,
+    create_rdma_connection,
+    create_socket_connection,
+)
 from exo.shared.topology import Topology
 from exo.shared.types.commands import PlaceInstance
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.events import InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.profiling import NetworkInterfaceInfo, NodePerformanceProfile
-from exo.shared.types.topology import Connection, NodeInfo
+from exo.shared.types.multiaddr import Multiaddr
+from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
+from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.instances import (
    Instance,
    InstanceId,
@@ -26,11 +30,6 @@ from exo.shared.types.worker.runners import ShardAssignments
 from exo.shared.types.worker.shards import Sharding


-@pytest.fixture
-def topology() -> Topology:
-    return Topology()
-
-
@pytest.fixture
 def instance() -> Instance:
    return MlxRingInstance(
@@ -77,34 +76,62 @@ def test_get_instance_placements_create_instance(
    available_memory: tuple[int, int, int],
    total_layers: int,
    expected_layers: tuple[int, int, int],
-    topology: Topology,
    model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
 ):
    # arrange
    model_meta.n_layers = total_layers
    model_meta.storage_size.in_bytes = sum(
        available_memory
    )  # make it exactly fit across all nodes
+    topology = Topology()

    cic = place_instance_command(model_meta)
    node_id_a = NodeId()
    node_id_b = NodeId()
    node_id_c = NodeId()
-    topology.add_node(create_node(available_memory[0], node_id_a))
-    topology.add_node(create_node(available_memory[1], node_id_b))
-    topology.add_node(create_node(available_memory[2], node_id_c))
-    # Add bidirectional connections for ring topology
-    topology.add_connection(create_connection(node_id_a, node_id_b))
-    topology.add_connection(create_connection(node_id_b, node_id_a))
-    topology.add_connection(create_connection(node_id_b, node_id_c))
-    topology.add_connection(create_connection(node_id_c, node_id_b))
-    topology.add_connection(create_connection(node_id_c, node_id_a))
-    topology.add_connection(create_connection(node_id_a, node_id_c))
+
+    # fully connected (directed) between the 3 nodes
+    conn_a_b = Connection(
+        source=node_id_a, sink=node_id_b, edge=create_socket_connection(1)
+    )
+    conn_b_c = Connection(
+        source=node_id_b, sink=node_id_c, edge=create_socket_connection(2)
+    )
+    conn_c_a = Connection(
+        source=node_id_c, sink=node_id_a, edge=create_socket_connection(3)
+    )
+    conn_c_b = Connection(
+        source=node_id_c, sink=node_id_b, edge=create_socket_connection(4)
+    )
+    conn_a_c = Connection(
+        source=node_id_a, sink=node_id_c, edge=create_socket_connection(5)
+    )
+    conn_b_a = Connection(
+        source=node_id_b, sink=node_id_a, edge=create_socket_connection(6)
+    )
+
+    node_memory = {
+        node_id_a: create_node_memory(available_memory[0]),
+        node_id_b: create_node_memory(available_memory[1]),
+        node_id_c: create_node_memory(available_memory[2]),
+    }
+    node_network = {
+        node_id_a: create_node_network(),
+        node_id_b: create_node_network(),
+        node_id_c: create_node_network(),
+    }
+    topology.add_node(node_id_a)
+    topology.add_node(node_id_b)
+    topology.add_node(node_id_c)
+    topology.add_connection(conn_a_b)
+    topology.add_connection(conn_b_c)
+    topology.add_connection(conn_c_a)
+    topology.add_connection(conn_c_b)
+    topology.add_connection(conn_a_c)
+    topology.add_connection(conn_b_a)

    # act
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, node_memory, node_network)

    # assert
    assert len(placements) == 1
@@ -130,12 +157,12 @@ def test_get_instance_placements_create_instance(
    assert shards_sorted[-1].end_layer == total_layers


-def test_get_instance_placements_one_node_exact_fit(
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-) -> None:
+def test_get_instance_placements_one_node_exact_fit() -> None:
    topology = Topology()
    node_id = NodeId()
-    topology.add_node(create_node(1000 * 1024, node_id))
+    topology.add_node(node_id)
+    node_memory = {node_id: create_node_memory(1000 * 1024)}
+    node_network = {node_id: create_node_network()}
    cic = place_instance_command(
        ModelMetadata(
            model_id=ModelId("test-model"),
@@ -146,7 +173,7 @@ def test_get_instance_placements_one_node_exact_fit(
            supports_tensor=True,
        ),
    )
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, node_memory, node_network)

    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
@@ -157,12 +184,12 @@ def test_get_instance_placements_one_node_exact_fit(
    assert len(instance.shard_assignments.runner_to_shard) == 1


-def test_get_instance_placements_one_node_fits_with_extra_memory(
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-) -> None:
+def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
    topology = Topology()
    node_id = NodeId()
-    topology.add_node(create_node(1001 * 1024, node_id))
+    topology.add_node(node_id)
+    node_memory = {node_id: create_node_memory(1001 * 1024)}
+    node_network = {node_id: create_node_network()}
    cic = place_instance_command(
        ModelMetadata(
            model_id=ModelId("test-model"),
@@ -173,7 +200,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory(
            supports_tensor=True,
        ),
    )
-    placements = place_instance(cic, topology, {})
+    placements = place_instance(cic, topology, {}, node_memory, node_network)

    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
@@ -184,12 +211,12 @@ def test_get_instance_placements_one_node_fits_with_extra_memory(
    assert len(instance.shard_assignments.runner_to_shard) == 1


-def test_get_instance_placements_one_node_not_fit(
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-) -> None:
+def test_get_instance_placements_one_node_not_fit() -> None:
    topology = Topology()
    node_id = NodeId()
-    topology.add_node(create_node(1000 * 1024, node_id))
+    topology.add_node(node_id)
+    node_memory = {node_id: create_node_memory(1000 * 1024)}
+    node_network = {node_id: create_node_network()}
    cic = place_instance_command(
        model_meta=ModelMetadata(
            model_id=ModelId("test-model"),
@@ -202,7 +229,7 @@ def test_get_instance_placements_one_node_not_fit(
    )

    with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
-        place_instance(cic, topology, {})
+        place_instance(cic, topology, {}, node_memory, node_network)


 def test_get_transition_events_no_change(instance: Instance):
@@ -247,179 +274,138 @@ def test_get_transition_events_delete_instance(instance: Instance):
    assert events[0].instance_id == instance_id


-def test_placement_selects_cycle_with_most_memory(
-    topology: Topology,
+def test_placement_selects_leaf_nodes(
    model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
 ):
-    # Arrange two 3-node cycles with different total memory.
-    # With bidirectional connections for ring topology, both cycles have non-leaf nodes.
-    # The algorithm should select the cycle with the most available memory.
+    # arrange
+    topology = Topology()

-    # Model requires more than any single node but fits within a 3-node cycle
-    model_meta.storage_size.in_bytes = 1500
-    model_meta.n_layers = 12
+    model_meta.storage_size = Memory.from_bytes(1000)

-    # Create node ids
    node_id_a = NodeId()
    node_id_b = NodeId()
    node_id_c = NodeId()
    node_id_d = NodeId()
-    node_id_e = NodeId()
-    node_id_f = NodeId()

-    # A-B-C cycle total memory = 1600 (< D-E-F total)
-    topology.add_node(create_node(400, node_id_a))
-    topology.add_node(create_node(400, node_id_b))
-    topology.add_node(create_node(800, node_id_c))
+    node_memory = {
+        node_id_a: create_node_memory(500),
+        node_id_b: create_node_memory(600),
+        node_id_c: create_node_memory(600),
+        node_id_d: create_node_memory(500),
+    }
+    node_network = {
+        node_id_a: create_node_network(),
+        node_id_b: create_node_network(),
+        node_id_c: create_node_network(),
+        node_id_d: create_node_network(),
+    }

-    # D-E-F cycle total memory = 1800 (> A-B-C total)
-    topology.add_node(create_node(600, node_id_d))
-    topology.add_node(create_node(600, node_id_e))
-    topology.add_node(create_node(600, node_id_f))
+    topology.add_node(node_id_a)
+    topology.add_node(node_id_b)
+    topology.add_node(node_id_c)
+    topology.add_node(node_id_d)

-    # Build bidirectional cycles for ring topology
-    topology.add_connection(create_connection(node_id_a, node_id_b))
-    topology.add_connection(create_connection(node_id_b, node_id_a))
-    topology.add_connection(create_connection(node_id_b, node_id_c))
-    topology.add_connection(create_connection(node_id_c, node_id_b))
-    topology.add_connection(create_connection(node_id_c, node_id_a))
-    topology.add_connection(create_connection(node_id_a, node_id_c))
-
-    topology.add_connection(create_connection(node_id_d, node_id_e))
-    topology.add_connection(create_connection(node_id_e, node_id_d))
-    topology.add_connection(create_connection(node_id_e, node_id_f))
-    topology.add_connection(create_connection(node_id_f, node_id_e))
-    topology.add_connection(create_connection(node_id_f, node_id_d))
-    topology.add_connection(create_connection(node_id_d, node_id_f))
-
-    cic = place_instance_command(
-        model_meta=model_meta,
+    # Daisy chain topology (directed)
+    topology.add_connection(
+        Connection(source=node_id_a, sink=node_id_b, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_b, sink=node_id_a, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_b, sink=node_id_c, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_c, sink=node_id_b, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_c, sink=node_id_d, edge=create_socket_connection(1))
+    )
+    topology.add_connection(
+        Connection(source=node_id_d, sink=node_id_c, edge=create_socket_connection(1))
    )

-    # Act
-    placements = place_instance(cic, topology, {})
+    cic = place_instance_command(model_meta=model_meta)

-    # Assert: D-E-F cycle should be selected as it has more total memory
+    # act
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
+
+    # assert
    assert len(placements) == 1
-    instance_id = list(placements.keys())[0]
-    instance = placements[instance_id]
+    instance = list(placements.values())[0]

    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    less_memory_cycle_nodes = {node_id_a, node_id_b, node_id_c}
-    more_memory_cycle_nodes = {node_id_d, node_id_e, node_id_f}
-
-    assert more_memory_cycle_nodes.issubset(assigned_nodes)
-    assert assigned_nodes.isdisjoint(less_memory_cycle_nodes)
+    assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set(
+        (
+            node_id_c,
+            node_id_d,
+        )
+    )


 def test_tensor_rdma_backend_connectivity_matrix(
-    topology: Topology,
    model_meta: ModelMetadata,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
 ):
+    # arrange
+    topology = Topology()
    model_meta.n_layers = 12
    model_meta.storage_size.in_bytes = 1500

-    node_id_a = NodeId()
-    node_id_b = NodeId()
-    node_id_c = NodeId()
+    node_a = NodeId()
+    node_b = NodeId()
+    node_c = NodeId()

-    node_a = create_node(500, node_id_a)
-    node_b = create_node(500, node_id_b)
-    node_c = create_node(500, node_id_c)
+    node_memory = {
+        node_a: create_node_memory(500),
+        node_b: create_node_memory(500),
+        node_c: create_node_memory(500),
+    }

    ethernet_interface = NetworkInterfaceInfo(
        name="en0",
-        ip_address="192.168.1.100",
+        ip_address="10.0.0.1",
+    )
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
    )

-    assert node_a.node_profile is not None
-    assert node_b.node_profile is not None
-    assert node_c.node_profile is not None
-
-    conn_a_b = create_connection(node_id_a, node_id_b)
-    conn_b_c = create_connection(node_id_b, node_id_c)
-    conn_c_a = create_connection(node_id_c, node_id_a)
-
-    conn_b_a = create_connection(node_id_b, node_id_a)
-    conn_c_b = create_connection(node_id_c, node_id_b)
-    conn_a_c = create_connection(node_id_a, node_id_c)
-
-    assert conn_a_b.send_back_multiaddr is not None
-    assert conn_b_c.send_back_multiaddr is not None
-    assert conn_c_a.send_back_multiaddr is not None
-
-    assert conn_b_a.send_back_multiaddr is not None
-    assert conn_c_b.send_back_multiaddr is not None
-    assert conn_a_c.send_back_multiaddr is not None
-
-    node_a.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_a.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_c_a.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_b_a.send_back_multiaddr.ip_address,
-            ),
-            ethernet_interface,
-        ],
-        system=node_a.node_profile.system,
-    )
-    node_b.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_b.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_c_b.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_a_b.send_back_multiaddr.ip_address,
-            ),
-            ethernet_interface,
-        ],
-        system=node_b.node_profile.system,
-    )
-    node_c.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_c.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_a_c.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_b_c.send_back_multiaddr.ip_address,
-            ),
-            ethernet_interface,
-        ],
-        system=node_c.node_profile.system,
-    )
+    node_network = {
+        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
+    }

    topology.add_node(node_a)
    topology.add_node(node_b)
    topology.add_node(node_c)
-    topology.add_connection(conn_a_b)
-    topology.add_connection(conn_b_c)
-    topology.add_connection(conn_c_a)
-    topology.add_connection(conn_b_a)
-    topology.add_connection(conn_c_b)
-    topology.add_connection(conn_a_c)
+
+    # RDMA connections (directed)
+    topology.add_connection(
+        Connection(source=node_a, sink=node_b, edge=create_rdma_connection(3))
+    )
+    topology.add_connection(
+        Connection(source=node_b, sink=node_a, edge=create_rdma_connection(3))
+    )
+    topology.add_connection(
+        Connection(source=node_b, sink=node_c, edge=create_rdma_connection(4))
+    )
+    topology.add_connection(
+        Connection(source=node_c, sink=node_b, edge=create_rdma_connection(4))
+    )
+    topology.add_connection(
+        Connection(source=node_a, sink=node_c, edge=create_rdma_connection(5))
+    )
+    topology.add_connection(
+        Connection(source=node_c, sink=node_a, edge=create_rdma_connection(5))
+    )
+
+    # Ethernet connections (directed)
+    topology.add_connection(Connection(source=node_a, sink=node_b, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_b, sink=node_c, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_c, sink=node_a, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_a, sink=node_c, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_b, sink=node_a, edge=ethernet_conn))
+    topology.add_connection(Connection(source=node_c, sink=node_b, edge=ethernet_conn))

    cic = PlaceInstance(
        sharding=Sharding.Tensor,
@@ -429,35 +415,34 @@ def test_tensor_rdma_backend_connectivity_matrix(
        min_nodes=1,
    )

-    placements = place_instance(cic, topology, {})
+    # act
+    placements = place_instance(cic, topology, {}, node_memory, node_network)

+    # assert
    assert len(placements) == 1
    instance_id = list(placements.keys())[0]
    instance = placements[instance_id]

    assert isinstance(instance, MlxJacclInstance)

-    assert instance.ibv_devices is not None
+    assert instance.jaccl_devices is not None
    assert instance.jaccl_coordinators is not None

-    matrix = instance.ibv_devices
+    matrix = instance.jaccl_devices
    assert len(matrix) == 3
-
    for i in range(3):
        assert matrix[i][i] is None

    assigned_nodes = list(instance.shard_assignments.node_to_runner.keys())
    node_to_idx = {node_id: idx for idx, node_id in enumerate(assigned_nodes)}

-    idx_a = node_to_idx[node_id_a]
-    idx_b = node_to_idx[node_id_b]
-    idx_c = node_to_idx[node_id_c]
+    idx_a = node_to_idx[node_a]
+    idx_b = node_to_idx[node_b]
+    idx_c = node_to_idx[node_c]

-    logger.info(matrix)
-
-    assert matrix[idx_a][idx_b] == "rdma_en4"
-    assert matrix[idx_b][idx_c] == "rdma_en3"
-    assert matrix[idx_c][idx_a] == "rdma_en3"
+    assert matrix[idx_a][idx_b] == "rdma_en3"
+    assert matrix[idx_b][idx_c] == "rdma_en4"
+    assert matrix[idx_c][idx_a] == "rdma_en5"

    # Verify coordinators are set for all nodes
    assert len(instance.jaccl_coordinators) == 3
@@ -469,7 +454,5 @@ def test_tensor_rdma_backend_connectivity_matrix(
        if node_id == assigned_nodes[0]:
            assert coordinator.startswith("0.0.0.0:")
        else:
-            # Non-rank-0 nodes should have valid IP addresses (can be link-local)
            ip_part = coordinator.split(":")[0]
-            # Just verify it's a valid IP format
            assert len(ip_part.split(".")) == 4
--- a/src/exo/master/tests/test_placement_utils.py
+++ b/src/exo/master/tests/test_placement_utils.py
@@ -1,5 +1,3 @@
-from typing import Callable
-
 import pytest

 from exo.master.placement_utils import (
@@ -10,154 +8,175 @@ 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.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId, ModelMetadata
-from exo.shared.types.profiling import NetworkInterfaceInfo, NodePerformanceProfile
-from exo.shared.types.topology import Connection, NodeInfo
+from exo.shared.types.profiling import (
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
+)
+from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.shards import Sharding


-@pytest.fixture
-def topology() -> Topology:
-    topology = Topology()
-    return topology
-
-
-def test_filter_cycles_by_memory(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_filter_cycles_by_memory():
    # arrange
    node1_id = NodeId()
    node2_id = NodeId()
+    connection1 = Connection(
+        source=node1_id, sink=node2_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node2_id, sink=node1_id, edge=create_socket_connection(2)
+    )

-    node1 = create_node(1000 * 1024, node1_id)
-    node2 = create_node(1000 * 1024, node2_id)
-
-    topology.add_node(node1)
-    topology.add_node(node2)
-
-    connection1 = create_connection(node1_id, node2_id)
-    connection2 = create_connection(node2_id, node1_id)
+    node1_mem = create_node_memory(1000 * 1024)
+    node2_mem = create_node_memory(1000 * 1024)
+    node_memory = {node1_id: node1_mem, node2_id: node2_mem}

+    topology = Topology()
+    topology.add_node(node1_id)
+    topology.add_node(node2_id)
    topology.add_connection(connection1)
    topology.add_connection(connection2)

-    cycles = topology.get_cycles()
+    cycles = [c for c in topology.get_cycles() if len(c) != 1]
    assert len(cycles) == 1
    assert len(cycles[0]) == 2

    # act
-    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_bytes(1))
+    filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_bytes(1))

    # assert
    assert len(filtered_cycles) == 1
    assert len(filtered_cycles[0]) == 2
-    assert set(n.node_id for n in filtered_cycles[0]) == {node1_id, node2_id}
+    assert set(n for n in filtered_cycles[0]) == {node1_id, node2_id}


-def test_filter_cycles_by_insufficient_memory(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_filter_cycles_by_insufficient_memory():
    # arrange
    node1_id = NodeId()
    node2_id = NodeId()
+    connection1 = Connection(
+        source=node1_id, sink=node2_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node2_id, sink=node1_id, edge=create_socket_connection(2)
+    )

-    node1 = create_node(1000 * 1024, node1_id)
-    node2 = create_node(1000 * 1024, node2_id)
-
-    topology.add_node(node1)
-    topology.add_node(node2)
-
-    connection1 = create_connection(node1_id, node2_id)
-    connection2 = create_connection(node2_id, node1_id)
+    node1_mem = create_node_memory(1000 * 1024)
+    node2_mem = create_node_memory(1000 * 1024)
+    node_memory = {node1_id: node1_mem, node2_id: node2_mem}

+    topology = Topology()
+    topology.add_node(node1_id)
+    topology.add_node(node2_id)
    topology.add_connection(connection1)
    topology.add_connection(connection2)

    # act
    filtered_cycles = filter_cycles_by_memory(
-        topology.get_cycles(), Memory.from_kb(2001)
+        topology.get_cycles(), node_memory, Memory.from_kb(2001)
    )

    # assert
    assert len(filtered_cycles) == 0


-def test_filter_multiple_cycles_by_memory(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_filter_multiple_cycles_by_memory():
    # arrange
    node_a_id = NodeId()
    node_b_id = NodeId()
    node_c_id = NodeId()
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_a_id, sink=node_c_id, edge=create_socket_connection(3)
+    )
+    connection4 = Connection(
+        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
+    )

-    node_a = create_node(500 * 1024, node_a_id)
-    node_b = create_node(500 * 1024, node_b_id)
-    node_c = create_node(1000 * 1024, node_c_id)
+    node_a_mem = create_node_memory(500 * 1024)
+    node_b_mem = create_node_memory(500 * 1024)
+    node_c_mem = create_node_memory(1000 * 1024)
+    node_memory = {
+        node_a_id: node_a_mem,
+        node_b_id: node_b_mem,
+        node_c_id: node_c_mem,
+    }

-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
-
-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
-
-    topology.add_connection(create_connection(node_a_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_b_id))
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
+    topology.add_connection(connection4)

    cycles = topology.get_cycles()

    # act
-    filtered_cycles = filter_cycles_by_memory(cycles, Memory.from_kb(1500))
+    filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_kb(1500))

    # assert
    assert len(filtered_cycles) == 1
    assert len(filtered_cycles[0]) == 3
-    assert set(n.node_id for n in filtered_cycles[0]) == {
+    assert set(n for n in filtered_cycles[0]) == {
        node_a_id,
        node_b_id,
        node_c_id,
    }


-def test_get_smallest_cycles(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_get_smallest_cycles():
    # arrange
    node_a_id = NodeId()
    node_b_id = NodeId()
    node_c_id = NodeId()

-    node_a = create_node(500 * 1024, node_a_id)
-    node_b = create_node(500 * 1024, node_b_id)
-    node_c = create_node(1000 * 1024, node_c_id)
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)

-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_a_id, sink=node_c_id, edge=create_socket_connection(3)
+    )
+    connection4 = Connection(
+        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
+    )

-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
+    topology.add_connection(connection4)
+
+    cycles = [c for c in topology.get_cycles() if len(c) != 1]  # ignore singletons

    # act
-    smallest_cycles = get_smallest_cycles(topology.get_cycles())
+    smallest_cycles = get_smallest_cycles(cycles)

    # assert
    assert len(smallest_cycles) == 1
    assert len(smallest_cycles[0]) == 2
-    assert set(n.node_id for n in smallest_cycles[0]) == {node_a_id, node_b_id}
+    assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id}


@pytest.mark.parametrize(
@@ -172,9 +191,6 @@ def test_get_smallest_cycles(
    ],
 )
 def test_get_shard_assignments(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
    available_memory: tuple[int, int, int],
    total_layers: int,
    expected_layers: tuple[int, int, int],
@@ -184,18 +200,37 @@ def test_get_shard_assignments(
    node_b_id = NodeId()
    node_c_id = NodeId()

-    node_a = create_node(available_memory[0] * 1024, node_a_id)
-    node_b = create_node(available_memory[1] * 1024, node_b_id)
-    node_c = create_node(available_memory[2] * 1024, node_c_id)
+    # create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists)
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
+    connection4 = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(4)
+    )

-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)
+    topology.add_connection(connection4)

-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
+    node_a_mem = create_node_memory(available_memory[0] * 1024)
+    node_b_mem = create_node_memory(available_memory[1] * 1024)
+    node_c_mem = create_node_memory(available_memory[2] * 1024)
+    node_memory = {
+        node_a_id: node_a_mem,
+        node_b_id: node_b_mem,
+        node_c_id: node_c_mem,
+    }

    model_meta = ModelMetadata(
        model_id=ModelId("test-model"),
@@ -205,23 +240,22 @@ def test_get_shard_assignments(
        hidden_size=1000,
        supports_tensor=True,
    )
+
    cycles = topology.get_cycles()
-    selected_cycle = cycles[0]
+
+    # pick the 3-node cycle deterministically (cycle ordering can vary)
+    selected_cycle = next(cycle for cycle in cycles if len(cycle) == 3)

    # act
    shard_assignments = get_shard_assignments(
-        model_meta, selected_cycle, Sharding.Pipeline
+        model_meta, selected_cycle, Sharding.Pipeline, node_memory=node_memory
    )

    # assert
    runner_id_a = shard_assignments.node_to_runner[node_a_id]
    runner_id_b = shard_assignments.node_to_runner[node_b_id]
    runner_id_c = shard_assignments.node_to_runner[node_c_id]
-    assert (
-        shard_assignments.runner_to_shard[runner_id_c].end_layer
-        - shard_assignments.runner_to_shard[runner_id_c].start_layer
-        == expected_layers[2]
-    )
+
    assert (
        shard_assignments.runner_to_shard[runner_id_a].end_layer
        - shard_assignments.runner_to_shard[runner_id_a].start_layer
@@ -232,30 +266,37 @@ def test_get_shard_assignments(
        - shard_assignments.runner_to_shard[runner_id_b].start_layer
        == expected_layers[1]
    )
+    assert (
+        shard_assignments.runner_to_shard[runner_id_c].end_layer
+        - shard_assignments.runner_to_shard[runner_id_c].start_layer
+        == expected_layers[2]
+    )


-def test_get_hosts_from_subgraph(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId, int | None], Connection],
-):
+def test_get_hosts_from_subgraph():
    # arrange
    node_a_id = NodeId()
    node_b_id = NodeId()
    node_c_id = NodeId()
+    topology = Topology()

-    node_a = create_node(500, node_a_id)
-    node_b = create_node(500, node_b_id)
-    node_c = create_node(1000, node_c_id)
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)

-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    connection1 = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    connection2 = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
+    )
+    connection3 = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
+    )

-    topology.add_connection(create_connection(node_a_id, node_b_id, 5001))
-    topology.add_connection(create_connection(node_b_id, node_c_id, 5002))
-    topology.add_connection(create_connection(node_c_id, node_a_id, 5003))
-    topology.add_connection(create_connection(node_b_id, node_a_id, 5004))
+    topology.add_connection(connection1)
+    topology.add_connection(connection2)
+    topology.add_connection(connection3)

    # act
    hosts = get_hosts_from_subgraph(topology)
@@ -263,95 +304,68 @@ def test_get_hosts_from_subgraph(
    # assert
    assert len(hosts) == 3
    expected_hosts = [
-        Host(ip=("169.254.0.2"), port=5001),
-        Host(ip=("169.254.0.3"), port=5002),
-        Host(ip=("169.254.0.4"), port=5003),
+        Host(ip="169.254.0.1", port=1234),
+        Host(ip="169.254.0.2", port=1234),
+        Host(ip="169.254.0.3", port=1234),
    ]
    for expected_host in expected_hosts:
        assert expected_host in hosts


-def test_get_mlx_jaccl_coordinators(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId, int | None], Connection],
-):
+def test_get_mlx_jaccl_coordinators():
    # arrange
    node_a_id = NodeId()
    node_b_id = NodeId()
    node_c_id = NodeId()

-    node_a = create_node(500 * 1024, node_a_id)
-    node_b = create_node(500 * 1024, node_b_id)
-    node_c = create_node(1000 * 1024, node_c_id)
-
-    conn_a_b = create_connection(node_a_id, node_b_id, 5001)
-    conn_b_a = create_connection(node_b_id, node_a_id, 5002)
-    conn_b_c = create_connection(node_b_id, node_c_id, 5003)
-    conn_c_b = create_connection(node_c_id, node_b_id, 5004)
-    conn_c_a = create_connection(node_c_id, node_a_id, 5005)
-    conn_a_c = create_connection(node_a_id, node_c_id, 5006)
-
-    # Update node profiles with network interfaces before adding to topology
-    assert node_a.node_profile is not None
-    assert node_b.node_profile is not None
-    assert node_c.node_profile is not None
-
-    node_a.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_a.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_a_b.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_a_c.send_back_multiaddr.ip_address,
-            ),
-        ],
-        system=node_a.node_profile.system,
+    # fully connected (directed) between the 3 nodes
+    conn_a_b = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
    )
-    node_b.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_b.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_b_a.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_b_c.send_back_multiaddr.ip_address,
-            ),
-        ],
-        system=node_b.node_profile.system,
+    conn_b_a = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(2)
    )
-    node_c.node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=node_c.node_profile.memory,
-        network_interfaces=[
-            NetworkInterfaceInfo(
-                name="en3",
-                ip_address=conn_c_b.send_back_multiaddr.ip_address,
-            ),
-            NetworkInterfaceInfo(
-                name="en4",
-                ip_address=conn_c_a.send_back_multiaddr.ip_address,
-            ),
-        ],
-        system=node_c.node_profile.system,
+    conn_b_c = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(3)
+    )
+    conn_c_b = Connection(
+        source=node_c_id, sink=node_b_id, edge=create_socket_connection(4)
+    )
+    conn_c_a = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(5)
+    )
+    conn_a_c = Connection(
+        source=node_a_id, sink=node_c_id, edge=create_socket_connection(6)
    )

-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    network_a = NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.5"),
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.2"),
+        ]
+    )
+    network_b = NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.1"),
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.4"),
+        ]
+    )
+    network_c = NodeNetworkInfo(
+        interfaces=[
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.3"),
+            NetworkInterfaceInfo(name="en0", ip_address="169.254.0.6"),
+        ]
+    )
+    node_network = {
+        node_a_id: network_a,
+        node_b_id: network_b,
+        node_c_id: network_c,
+    }
+
+    topology = Topology()
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)

    topology.add_connection(conn_a_b)
    topology.add_connection(conn_b_a)
@@ -360,11 +374,12 @@ def test_get_mlx_jaccl_coordinators(
    topology.add_connection(conn_c_a)
    topology.add_connection(conn_a_c)

-    cycle = [node_a, node_b, node_c]
-
    # act
    coordinators = get_mlx_jaccl_coordinators(
-        cycle, coordinator_port=5000, cycle_digraph=topology
+        node_a_id,
+        coordinator_port=5000,
+        cycle_digraph=topology,
+        node_network=node_network,
    )

    # assert
@@ -385,21 +400,22 @@ def test_get_mlx_jaccl_coordinators(
            f"Coordinator for {node_id} should use port 5000"
        )

-    # Rank 0 (node_a) treats this as the listen socket so should listen on all
-    # IPs
+    # Rank 0 (node_a) treats this as the listen socket so should listen on all IPs
    assert coordinators[node_a_id].startswith("0.0.0.0:"), (
-        "Rank 0 node should use localhost as coordinator"
+        "Rank 0 node should use 0.0.0.0 as coordinator listen address"
    )

    # Non-rank-0 nodes should use the specific IP from their connection to rank 0
    # node_b uses the IP from conn_b_a (node_b -> node_a)
-    assert coordinators[node_b_id] == (
-        f"{conn_b_a.send_back_multiaddr.ip_address}:5000"
+    assert isinstance(conn_b_a.edge, SocketConnection)
+    assert (
+        coordinators[node_b_id] == f"{conn_b_a.edge.sink_multiaddr.ip_address}:5000"
    ), "node_b should use the IP from conn_b_a"

    # node_c uses the IP from conn_c_a (node_c -> node_a)
+    assert isinstance(conn_c_a.edge, SocketConnection)
    assert coordinators[node_c_id] == (
-        f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
+        f"{conn_c_a.edge.sink_multiaddr.ip_address}:5000"
    ), "node_c should use the IP from conn_c_a"


@@ -457,29 +473,44 @@ class TestAllocateLayersProportionally:
        assert sum(result) == 3


-def test_get_shard_assignments_insufficient_memory_raises(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
+def test_get_shard_assignments_insufficient_memory_raises():
    """Test that ValueError is raised when a node has insufficient memory for its layers."""
    node_a_id = NodeId()
    node_b_id = NodeId()
    node_c_id = NodeId()
+    topology = Topology()

    # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers)
-    node_a = 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
+    node_a_mem = create_node_memory(900 * 1024)
+    node_b_mem = create_node_memory(50 * 1024)
+    node_c_mem = create_node_memory(10 * 1024)  # Insufficient memory

-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
+    topology.add_node(node_a_id)
+    topology.add_node(node_b_id)
+    topology.add_node(node_c_id)

-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
+    conn_a_b = Connection(
+        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
+    )
+    conn_b_c = Connection(
+        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
+    )
+    conn_c_a = Connection(
+        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
+    conn_b_a = Connection(
+        source=node_b_id, sink=node_a_id, edge=create_socket_connection(3)
+    )
+    topology.add_connection(conn_a_b)
+    topology.add_connection(conn_b_c)
+    topology.add_connection(conn_c_a)
+    topology.add_connection(conn_b_a)
+
+    node_memory = {
+        node_a_id: node_a_mem,
+        node_b_id: node_b_mem,
+        node_c_id: node_c_mem,
+    }

    model_meta = ModelMetadata(
        model_id=ModelId("test-model"),
@@ -493,4 +524,6 @@ def test_get_shard_assignments_insufficient_memory_raises(
    selected_cycle = cycles[0]

    with pytest.raises(ValueError, match="insufficient memory"):
-        get_shard_assignments(model_meta, selected_cycle, Sharding.Pipeline)
+        get_shard_assignments(
+            model_meta, selected_cycle, Sharding.Pipeline, node_memory
+        )
--- a/src/exo/master/tests/test_topology.py
+++ b/src/exo/master/tests/test_topology.py
@@ -1,13 +1,9 @@
 import pytest

 from exo.shared.topology import Topology
+from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
-)
-from exo.shared.types.topology import Connection, ConnectionProfile, NodeId, NodeInfo
+from exo.shared.types.topology import Connection, SocketConnection


@pytest.fixture
@@ -16,189 +12,97 @@ def topology() -> Topology:


@pytest.fixture
-def connection() -> Connection:
-    return Connection(
-        local_node_id=NodeId(),
-        send_back_node_id=NodeId(),
-        send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
-        connection_profile=ConnectionProfile(
-            throughput=1000, latency=1000, jitter=1000
-        ),
+def socket_connection() -> SocketConnection:
+    return SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
    )


-@pytest.fixture
-def node_profile() -> NodePerformanceProfile:
-    memory_profile = MemoryPerformanceProfile.from_bytes(
-        ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
-    )
-    system_profile = SystemPerformanceProfile()
-    return NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=memory_profile,
-        network_interfaces=[],
-        system=system_profile,
-    )
-
-
-@pytest.fixture
-def connection_profile() -> ConnectionProfile:
-    return ConnectionProfile(throughput=1000, latency=1000, jitter=1000)
-
-
-def test_add_node(topology: Topology, node_profile: NodePerformanceProfile):
+def test_add_node(topology: Topology):
    # arrange
    node_id = NodeId()

    # act
-    topology.add_node(NodeInfo(node_id=node_id, node_profile=node_profile))
+    topology.add_node(node_id)

    # assert
-    data = topology.get_node_profile(node_id)
-    assert data == node_profile
+    assert topology.node_is_leaf(node_id)


-def test_add_connection(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
+def test_add_connection(topology: Topology, socket_connection: SocketConnection):
    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
+    node_a = NodeId()
+    node_b = NodeId()
+    connection = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
    topology.add_connection(connection)

    # act
-    data = topology.get_connection_profile(connection)
+    data = list(topology.list_connections())

    # assert
-    assert data == connection.connection_profile
+    assert data == [connection]

-
-def test_update_node_profile(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
-    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
-
-    new_node_profile = NodePerformanceProfile(
-        model_id="test",
-        chip_id="test",
-        friendly_name="test",
-        memory=MemoryPerformanceProfile.from_bytes(
-            ram_total=1000, ram_available=1000, swap_total=1000, swap_available=1000
-        ),
-        network_interfaces=[],
-        system=SystemPerformanceProfile(),
-    )
-
-    # act
-    topology.update_node_profile(
-        connection.local_node_id, node_profile=new_node_profile
-    )
-
-    # assert
-    data = topology.get_node_profile(connection.local_node_id)
-    assert data == new_node_profile
-
-
-def test_update_connection_profile(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
-    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
-
-    new_connection_profile = ConnectionProfile(
-        throughput=2000, latency=2000, jitter=2000
-    )
-    connection = Connection(
-        local_node_id=connection.local_node_id,
-        send_back_node_id=connection.send_back_node_id,
-        send_back_multiaddr=connection.send_back_multiaddr,
-        connection_profile=new_connection_profile,
-    )
-
-    # act
-    topology.update_connection_profile(connection)
-
-    # assert
-    data = topology.get_connection_profile(connection)
-    assert data == new_connection_profile
+    assert topology.node_is_leaf(node_a)
+    assert topology.node_is_leaf(node_b)


 def test_remove_connection_still_connected(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+    topology: Topology, socket_connection: SocketConnection
 ):
    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
+    node_a = NodeId()
+    node_b = NodeId()
+    conn = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(conn)

    # act
-    topology.remove_connection(connection)
+    topology.remove_connection(conn)

    # assert
-    assert topology.get_connection_profile(connection) is None
+    assert list(topology.get_all_connections_between(node_a, node_b)) == []


 def test_remove_node_still_connected(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
+    topology: Topology, socket_connection: SocketConnection
 ):
    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
+    node_a = NodeId()
+    node_b = NodeId()
+    conn = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(conn)
+    assert list(topology.out_edges(node_a)) == [conn]

    # act
-    topology.remove_node(connection.local_node_id)
+    topology.remove_node(node_b)

    # assert
-    assert topology.get_node_profile(connection.local_node_id) is None
+    assert list(topology.out_edges(node_a)) == []


-def test_list_nodes(
-    topology: Topology, node_profile: NodePerformanceProfile, connection: Connection
-):
+def test_list_nodes(topology: Topology, socket_connection: SocketConnection):
    # arrange
-    topology.add_node(
-        NodeInfo(node_id=connection.local_node_id, node_profile=node_profile)
-    )
-    topology.add_node(
-        NodeInfo(node_id=connection.send_back_node_id, node_profile=node_profile)
-    )
-    topology.add_connection(connection)
+    node_a = NodeId()
+    node_b = NodeId()
+    conn = Connection(source=node_a, sink=node_b, edge=socket_connection)
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(conn)
+    assert list(topology.out_edges(node_a)) == [conn]

    # act
    nodes = list(topology.list_nodes())

    # assert
    assert len(nodes) == 2
-    assert all(isinstance(node, NodeInfo) for node in nodes)
-    assert {node.node_id for node in nodes} == {
-        connection.local_node_id,
-        connection.send_back_node_id,
-    }
+    assert all(isinstance(node, NodeId) for node in nodes)
+    assert set(node for node in nodes) == set([node_a, node_b])
--- a/src/exo/plugins/init.py
+++ b/src/exo/plugins/init.py
@@ -0,0 +1,16 @@
+"""Exo Plugin System.
+
+This module provides the plugin architecture for extending exo with custom
+workload types (simulations, ML frameworks, etc.) without modifying core code.
+"""
+
+from exo.plugins.base import ExoPlugin, PluginCommand, PluginInstance
+from exo.plugins.registry import PluginRegistry, discover_plugins
+
+__all__ = [
+    "ExoPlugin",
+    "PluginCommand",
+    "PluginInstance",
+    "PluginRegistry",
+    "discover_plugins",
+]
--- a/src/exo/plugins/base.py
+++ b/src/exo/plugins/base.py
@@ -0,0 +1,171 @@
+"""Base classes and protocols for Exo plugins."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field
+
+from exo.shared.types.common import CommandId
+from exo.shared.types.events import Event
+from exo.shared.types.tasks import Task
+from exo.shared.types.worker.instances import InstanceId
+from exo.shared.types.worker.runners import RunnerId
+from exo.utils.pydantic_ext import TaggedModel
+
+if TYPE_CHECKING:
+    from exo.shared.topology import Topology
+    from exo.shared.types.worker.instances import BoundInstance, Instance
+    from exo.utils.channels import MpReceiver, MpSender
+    from exo.worker.runner.runner_supervisor import RunnerSupervisor
+
+
+class PluginCommand(TaggedModel):
+    """Base class for plugin-defined commands.
+
+    All plugin commands must inherit from this class. Commands are serialized
+    with their class name as a tag for routing.
+    """
+
+    command_id: CommandId = Field(default_factory=CommandId)
+
+
+class PluginInstance(TaggedModel):
+    """Base class for plugin-defined instances.
+
+    All plugin instances must inherit from this class. Plugins are expected
+    to define their own instance type with workload-specific fields.
+    """
+
+    instance_id: InstanceId
+
+
+class ExoPlugin(ABC):
+    """Protocol that all exo plugins must implement.
+
+    A plugin provides:
+    - Custom command types for API -> Master communication
+    - Custom instance types representing running workloads
+    - Placement logic for distributing work across nodes
+    - Planning logic for local task scheduling
+    - Runner implementation for executing work
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Unique identifier for this plugin (e.g., 'flash', 'pytorch', 'mpi')."""
+        ...
+
+    @property
+    @abstractmethod
+    def version(self) -> str:
+        """Semantic version string (e.g., '1.0.0')."""
+        ...
+
+    # ========== Type Registration ==========
+
+    @abstractmethod
+    def get_command_types(self) -> Sequence[type]:
+        """Return command types this plugin handles.
+
+        These commands are routed to this plugin's process_command method.
+        Can return core BaseCommand types or PluginCommand types.
+        """
+        ...
+
+    @abstractmethod
+    def get_instance_type(self) -> type:
+        """Return the instance type this plugin creates.
+
+        This instance type is used for routing in planning and runner bootstrap.
+        Can return core Instance types or PluginInstance types.
+        """
+        ...
+
+    # ========== API Routes ==========
+
+    @abstractmethod
+    def get_api_routes(
+        self,
+    ) -> Sequence[tuple[str, str, Callable[..., Any]]]:
+        """Return FastAPI routes to register.
+
+        Each tuple: (method, path, handler)
+        Example: [('post', '/flash/launch', self.launch_handler)]
+
+        Handlers receive a PluginContext with access to:
+        - state: Current State object
+        - send_command: Async function to send commands
+        - node_id: Current node's ID
+        """
+        ...
+
+    # ========== Master Command Handling ==========
+
+    @abstractmethod
+    def handles_command(self, command: Any) -> bool:  # pyright: ignore[reportAny]
+        """Return True if this plugin handles the given command type."""
+        ...
+
+    @abstractmethod
+    def process_command(
+        self,
+        command: Any,  # pyright: ignore[reportAny]
+        topology: "Topology",
+        current_instances: Mapping[InstanceId, "Instance"],
+    ) -> Sequence[Event]:
+        """Process a command and return events to emit.
+
+        Typically creates placement and returns InstanceCreated/InstanceDeleted events.
+
+        Args:
+            command: The command to process
+            topology: Current cluster topology
+            current_instances: Currently running instances
+
+        Returns:
+            Sequence of events to emit (e.g., InstanceCreated, InstanceDeleted)
+        """
+        ...
+
+    # ========== Worker Planning ==========
+
+    @abstractmethod
+    def handles_instance(self, instance: object) -> bool:
+        """Return True if this plugin manages the given instance type."""
+        ...
+
+    @abstractmethod
+    def plan_task(
+        self,
+        runners: Mapping[RunnerId, "RunnerSupervisor"],
+        instances: Mapping[InstanceId, "Instance"],
+    ) -> Task | None:
+        """Plan the next task for plugin instances.
+
+        Called during each planning cycle.
+        Return None if no task is needed.
+        """
+        ...
+
+    @abstractmethod
+    def should_skip_download(self, instance: object) -> bool:
+        """Return True if this instance type doesn't need model downloads."""
+        ...
+
+    # ========== Runner Bootstrap ==========
+
+    @abstractmethod
+    def create_runner(
+        self,
+        bound_instance: "BoundInstance",
+        event_sender: "MpSender[Event]",
+        task_receiver: "MpReceiver[Task]",
+    ) -> None:
+        """Entry point for the runner process.
+
+        Called in a subprocess to execute the actual workload.
+        This function should block until the workload completes.
+        """
+        ...
--- a/src/exo/plugins/context.py
+++ b/src/exo/plugins/context.py
@@ -0,0 +1,21 @@
+"""Context objects passed to plugin handlers."""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from exo.shared.types.commands import Command
+from exo.shared.types.common import NodeId
+from exo.shared.types.state import State
+
+
+@dataclass
+class PluginContext:
+    """Context provided to plugin API handlers.
+
+    This gives plugins access to the current state and the ability to send
+    commands without direct access to internal API components.
+    """
+
+    state: State
+    send_command: Callable[[Command], Awaitable[None]]
+    node_id: NodeId
--- a/src/exo/plugins/implementations/init.py
+++ b/src/exo/plugins/implementations/init.py
@@ -0,0 +1,5 @@
+"""Plugin implementations directory.
+
+Each subdirectory should contain a plugin with a register() function
+that returns an ExoPlugin instance.
+"""
--- a/src/exo/plugins/implementations/flash/init.py
+++ b/src/exo/plugins/implementations/flash/init.py
@@ -0,0 +1,8 @@
+"""FLASH Plugin - MPI-based simulation support for Exo."""
+
+from exo.plugins.implementations.flash.plugin import FLASHPlugin
+
+
+def register() -> FLASHPlugin:
+    """Entry point for plugin discovery."""
+    return FLASHPlugin()
--- a/src/exo/plugins/implementations/flash/api_handlers.py
+++ b/src/exo/plugins/implementations/flash/api_handlers.py
@@ -0,0 +1,108 @@
+"""FLASH plugin API handlers."""
+
+from typing import Any
+
+from fastapi import HTTPException
+
+from exo.plugins.context import PluginContext
+
+# Use core types for serialization compatibility
+from exo.shared.types.commands import LaunchFLASH, StopFLASH
+from exo.shared.types.worker.instances import FLASHInstance
+
+
+async def handle_launch_flash(
+    ctx: PluginContext,
+    simulation_name: str,
+    flash_executable_path: str,
+    working_directory: str,
+    parameter_file_path: str = "",
+    ranks_per_node: int = 1,
+    min_nodes: int = 1,
+    hosts: str = "",
+) -> dict[str, str]:
+    """Launch a FLASH MPI simulation across the cluster.
+
+    Args:
+        ctx: Plugin context with state and send_command
+        simulation_name: Name of the simulation
+        flash_executable_path: Path to the FLASH executable
+        working_directory: Working directory for the simulation
+        parameter_file_path: Path to parameter file (optional)
+        ranks_per_node: Number of MPI ranks per node
+        min_nodes: Minimum number of nodes required
+        hosts: Optional comma-separated hostnames (e.g., "s14,james21-1").
+               If not provided, IPs are discovered from topology edges.
+    """
+    command = LaunchFLASH(
+        simulation_name=simulation_name,
+        flash_executable_path=flash_executable_path,
+        parameter_file_path=parameter_file_path,
+        working_directory=working_directory,
+        ranks_per_node=ranks_per_node,
+        min_nodes=min_nodes,
+        hosts=hosts,
+    )
+    await ctx.send_command(command)
+
+    return {
+        "message": "FLASH launch command received",
+        "command_id": str(command.command_id),
+        "simulation_name": simulation_name,
+    }
+
+
+async def handle_stop_flash(
+    ctx: PluginContext,
+    instance_id: str,
+) -> dict[str, str]:
+    """Stop a running FLASH simulation."""
+    from exo.shared.types.worker.instances import InstanceId
+
+    inst_id = InstanceId(instance_id)
+
+    if inst_id not in ctx.state.instances:
+        raise HTTPException(status_code=404, detail="Instance not found")
+
+    instance = ctx.state.instances[inst_id]
+    if not isinstance(instance, FLASHInstance):
+        raise HTTPException(
+            status_code=400, detail="Instance is not a FLASH simulation"
+        )
+
+    command = StopFLASH(instance_id=inst_id)
+    await ctx.send_command(command)
+
+    return {
+        "message": "Stop command received",
+        "command_id": str(command.command_id),
+        "instance_id": str(instance_id),
+    }
+
+
+async def handle_list_flash_instances(ctx: PluginContext) -> list[dict[str, Any]]:
+    """List all FLASH simulation instances."""
+    flash_instances: list[dict[str, Any]] = []
+    for instance_id, instance in ctx.state.instances.items():
+        if isinstance(instance, FLASHInstance):
+            # Get runner statuses for this instance
+            runner_statuses: dict[str, str | None] = {}
+            for (
+                node_id,
+                runner_id,
+            ) in instance.shard_assignments.node_to_runner.items():
+                runner_status = ctx.state.runners.get(runner_id)
+                runner_statuses[str(node_id)] = (
+                    str(runner_status) if runner_status else None
+                )
+
+            flash_instances.append(
+                {
+                    "instance_id": str(instance_id),
+                    "simulation_name": instance.simulation_name,
+                    "total_ranks": instance.total_ranks,
+                    "working_directory": instance.working_directory,
+                    "runner_statuses": runner_statuses,
+                }
+            )
+    return flash_instances
--- a/src/exo/plugins/implementations/flash/placement.py
+++ b/src/exo/plugins/implementations/flash/placement.py
@@ -0,0 +1,152 @@
+"""FLASH plugin placement logic."""
+
+from collections.abc import Mapping
+from copy import deepcopy
+
+from loguru import logger
+
+from exo.shared.topology import Topology
+from exo.shared.types.commands import LaunchFLASH
+from exo.shared.types.common import Host, NodeId
+from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.topology import SocketConnection
+from exo.shared.types.worker.instances import FLASHInstance, Instance, InstanceId
+from exo.shared.types.worker.runners import (
+    RunnerId,
+    ShardAssignments,
+)
+from exo.shared.types.worker.shards import PipelineShardMetadata
+
+
+def place_flash_instance(
+    command: LaunchFLASH,
+    topology: Topology,
+    current_instances: Mapping[InstanceId, Instance],
+) -> dict[InstanceId, Instance]:
+    """Place a FLASH simulation instance across available nodes.
+
+    Unlike MLX instances which use ring/JACCL topology for tensor parallelism,
+    FLASH instances use MPI for communication. We just need to provide the
+    node IPs so the runner can generate an MPI hostfile.
+    """
+    instance_id = InstanceId()
+    target_instances: dict[InstanceId, Instance] = dict(deepcopy(current_instances))
+
+    all_nodes = list(topology.list_nodes())
+
+    if len(all_nodes) < command.min_nodes:
+        raise ValueError(
+            f"Not enough nodes: need {command.min_nodes}, have {len(all_nodes)}"
+        )
+
+    # Select nodes (take the first min_nodes)
+    selected_nodes = all_nodes[: command.min_nodes]
+
+    logger.info(
+        f"Placing FLASH instance '{command.simulation_name}' on {len(selected_nodes)} nodes"
+    )
+
+    # Build shard assignments (one runner per node for FLASH)
+    runner_to_shard: dict[RunnerId, PipelineShardMetadata] = {}
+    node_to_runner: dict[NodeId, RunnerId] = {}
+
+    # Create a dummy ModelMetadata for FLASH (required by ShardMetadata interface)
+    flash_model_meta = ModelMetadata(
+        model_id=ModelId(command.simulation_name),
+        pretty_name=f"FLASH: {command.simulation_name}",
+        storage_size=Memory(in_bytes=0),
+        n_layers=1,
+        hidden_size=1,
+        supports_tensor=False,
+    )
+
+    for i, node_id in enumerate(selected_nodes):
+        runner_id = RunnerId()
+        node_to_runner[node_id] = runner_id
+        runner_to_shard[runner_id] = PipelineShardMetadata(
+            device_rank=i,
+            world_size=len(selected_nodes),
+            model_meta=flash_model_meta,
+            start_layer=0,
+            end_layer=1,
+            n_layers=1,
+        )
+
+    shard_assignments = ShardAssignments(
+        model_id=ModelId(command.simulation_name),
+        runner_to_shard=runner_to_shard,
+        node_to_runner=node_to_runner,
+    )
+
+    # Build hosts_by_node - get hostnames/IPs for MPI hostfile generation
+    hosts_by_node: dict[NodeId, list[Host]] = {}
+
+    # If explicit hosts are provided, use them directly
+    if command.hosts:
+        explicit_hosts = [h.strip() for h in command.hosts.split(",") if h.strip()]
+        logger.info(f"FLASH placement: explicit hosts provided: {explicit_hosts}")
+        for i, node_id in enumerate(selected_nodes):
+            if i < len(explicit_hosts):
+                hosts_by_node[node_id] = [Host(ip=explicit_hosts[i], port=0)]
+                logger.info(
+                    f"FLASH placement: node {node_id} (rank {i}) -> IP {explicit_hosts[i]}"
+                )
+            else:
+                logger.warning(
+                    f"Not enough hosts provided for node {i}, using localhost"
+                )
+                hosts_by_node[node_id] = [Host(ip="127.0.0.1", port=0)]
+        logger.info(
+            f"FLASH placement: coordinator will be rank 0 at IP {explicit_hosts[0]}"
+        )
+    else:
+        # Try to get IPs from topology edges
+        for node_id in selected_nodes:
+            node_hosts: list[Host] = []
+
+            # Get IP from outgoing edges (connections to other nodes via mDNS discovery)
+            for conn in topology.out_edges(node_id):
+                if isinstance(conn.edge, SocketConnection):
+                    # Extract IP from multiaddr
+                    ip = conn.edge.sink_multiaddr.ip_address
+                    # Skip link-local and localhost addresses
+                    if not ip.startswith("169.254.") and not ip.startswith("127."):
+                        node_hosts.append(Host(ip=ip, port=0))
+                        break
+
+            # Last resort: use localhost (will only work for single-node)
+            if not node_hosts:
+                logger.warning(
+                    f"Could not determine IP for node {node_id}, using localhost"
+                )
+                node_hosts.append(Host(ip="127.0.0.1", port=0))
+
+            hosts_by_node[node_id] = node_hosts
+
+    total_ranks = len(selected_nodes) * command.ranks_per_node
+
+    # Determine coordinator IP - first node's first host IP
+    first_node_id: NodeId = next(iter(hosts_by_node.keys()))
+    coordinator_ip: str = (
+        hosts_by_node[first_node_id][0].ip
+        if hosts_by_node[first_node_id]
+        else "127.0.0.1"
+    )
+
+    target_instances[instance_id] = FLASHInstance(
+        instance_id=instance_id,
+        shard_assignments=shard_assignments,
+        hosts_by_node=hosts_by_node,
+        flash_executable_path=command.flash_executable_path,
+        parameter_file_path=command.parameter_file_path,
+        working_directory=command.working_directory,
+        ranks_per_node=command.ranks_per_node,
+        total_ranks=total_ranks,
+        simulation_name=command.simulation_name,
+        coordinator_ip=coordinator_ip,
+    )
+
+    logger.info(f"Created FLASH instance {instance_id} with {total_ranks} total ranks")
+
+    return target_instances
--- a/src/exo/plugins/implementations/flash/planning.py
+++ b/src/exo/plugins/implementations/flash/planning.py
@@ -0,0 +1,36 @@
+"""FLASH plugin planning logic."""
+
+from collections.abc import Mapping
+
+from exo.shared.types.tasks import LoadModel, Task
+from exo.shared.types.worker.instances import FLASHInstance, Instance, InstanceId
+from exo.shared.types.worker.runners import RunnerId, RunnerIdle
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+
+
+def plan_flash(
+    runners: Mapping[RunnerId, RunnerSupervisor],
+    instances: Mapping[InstanceId, Instance],
+) -> Task | None:
+    """Plan tasks specifically for FLASH instances.
+
+    FLASH instances have a simpler lifecycle:
+    - CreateRunner (handled by core _create_runner)
+    - LoadModel (starts the simulation immediately)
+    - Shutdown (handled by core _kill_runner)
+
+    This function handles the LoadModel step for FLASH instances,
+    skipping the MLX-specific download/init/warmup steps.
+    """
+    for runner in runners.values():
+        instance = runner.bound_instance.instance
+
+        # Only handle FLASH instances
+        if not isinstance(instance, FLASHInstance):
+            continue
+
+        # If runner is idle, emit LoadModel to start the simulation
+        if isinstance(runner.status, RunnerIdle):
+            return LoadModel(instance_id=instance.instance_id)
+
+    return None
--- a/src/exo/plugins/implementations/flash/plugin.py
+++ b/src/exo/plugins/implementations/flash/plugin.py
@@ -0,0 +1,98 @@
+"""FLASH Plugin - Main plugin class."""
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+from exo.plugins.base import ExoPlugin
+from exo.plugins.implementations.flash.api_handlers import (
+    handle_launch_flash,
+    handle_list_flash_instances,
+    handle_stop_flash,
+)
+from exo.plugins.implementations.flash.placement import place_flash_instance
+from exo.plugins.implementations.flash.planning import plan_flash
+from exo.plugins.implementations.flash.runner import main as flash_runner_main
+from exo.shared.topology import Topology
+from exo.shared.types.commands import DeleteInstance, LaunchFLASH, StopFLASH
+from exo.shared.types.events import Event
+from exo.shared.types.tasks import Task
+from exo.shared.types.worker.instances import (
+    BoundInstance,
+    FLASHInstance,
+    Instance,
+    InstanceId,
+)
+from exo.shared.types.worker.runners import RunnerId
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+
+
+class FLASHPlugin(ExoPlugin):
+    """Plugin for FLASH MPI simulations."""
+
+    @property
+    def name(self) -> str:
+        return "flash"
+
+    @property
+    def version(self) -> str:
+        return "1.0.0"
+
+    def get_command_types(self) -> Sequence[type]:
+        return [LaunchFLASH, StopFLASH]
+
+    def get_instance_type(self) -> type:
+        return FLASHInstance
+
+    def get_api_routes(
+        self,
+    ) -> Sequence[tuple[str, str, Callable[..., Any]]]:
+        return [
+            ("post", "/flash/launch", handle_launch_flash),
+            ("delete", "/flash/{instance_id}", handle_stop_flash),
+            ("get", "/flash/instances", handle_list_flash_instances),
+        ]
+
+    def handles_command(self, command: Any) -> bool:  # pyright: ignore[reportAny]
+        return isinstance(command, (LaunchFLASH, StopFLASH))
+
+    def process_command(
+        self,
+        command: Any,  # pyright: ignore[reportAny]
+        topology: Topology,
+        current_instances: Mapping[InstanceId, Instance],
+    ) -> Sequence[Event]:
+        from exo.master.placement import delete_instance, get_transition_events
+
+        if isinstance(command, LaunchFLASH):
+            placement = place_flash_instance(command, topology, current_instances)
+            return list(get_transition_events(current_instances, placement))
+        elif isinstance(command, StopFLASH):
+            placement = delete_instance(
+                DeleteInstance(instance_id=command.instance_id),
+                current_instances,
+            )
+            return list(get_transition_events(current_instances, placement))
+        return []
+
+    def handles_instance(self, instance: object) -> bool:
+        return isinstance(instance, FLASHInstance)
+
+    def plan_task(
+        self,
+        runners: Mapping[RunnerId, RunnerSupervisor],
+        instances: Mapping[InstanceId, Instance],
+    ) -> Task | None:
+        return plan_flash(runners, instances)
+
+    def should_skip_download(self, instance: object) -> bool:
+        # FLASH instances don't need model downloads
+        return True
+
+    def create_runner(
+        self,
+        bound_instance: BoundInstance,
+        event_sender: MpSender[Event],
+        task_receiver: MpReceiver[Task],
+    ) -> None:
+        flash_runner_main(bound_instance, event_sender, task_receiver)
--- a/src/exo/plugins/implementations/flash/runner.py
+++ b/src/exo/plugins/implementations/flash/runner.py
@@ -0,0 +1,302 @@
+"""FLASH MPI Runner - spawns and monitors FLASH simulations.
+
+Exo-native distributed MPI:
+- Exo handles node discovery and coordination
+- Coordinator generates hostfile from Exo topology
+- mpirun uses exo-rsh (no SSH required) to spawn on remote nodes
+- exo-rsh connects to each node's Exo API (/execute endpoint) for remote execution
+- Workers just report ready and wait
+"""
+
+import os
+import shutil
+import socket
+import subprocess
+import threading
+
+from loguru import logger
+
+from exo.shared.types.events import (
+    Event,
+    RunnerStatusUpdated,
+    TaskAcknowledged,
+    TaskStatusUpdated,
+)
+from exo.shared.types.tasks import (
+    LoadModel,
+    Shutdown,
+    Task,
+    TaskStatus,
+)
+from exo.shared.types.worker.instances import BoundInstance, FLASHInstance
+from exo.shared.types.worker.runners import (
+    RunnerFailed,
+    RunnerIdle,
+    RunnerLoading,
+    RunnerReady,
+    RunnerRunning,
+    RunnerShutdown,
+    RunnerShuttingDown,
+    RunnerStatus,
+)
+from exo.utils.channels import MpReceiver, MpSender
+
+# Find mpirun in PATH, fallback to common locations
+MPIRUN_PATH = shutil.which("mpirun") or "/opt/homebrew/bin/mpirun"
+
+# exo-rsh is installed as console script by exo package
+_exo_rsh_path = shutil.which("exo-rsh")
+if not _exo_rsh_path:
+    raise RuntimeError("exo-rsh not found in PATH - this should be installed with exo")
+EXO_RSH_PATH: str = _exo_rsh_path
+
+
+def get_my_rank(instance: FLASHInstance, my_node_id: str) -> int:
+    """Determine this node's rank based on position in hosts_by_node."""
+    for i, node_id in enumerate(instance.hosts_by_node.keys()):
+        if str(node_id) == str(my_node_id):
+            return i
+    return -1
+
+
+def get_coordinator_host(instance: FLASHInstance) -> str:
+    """Get the IP of the coordinator node."""
+    return instance.coordinator_ip
+
+
+def resolve_host(host: str) -> str:
+    """Resolve host string to a usable hostname for MPI hostfile.
+
+    Accepts either an IP address or hostname. For IPs, attempts to resolve
+    to a hostname via DNS/mDNS. Hostnames are returned as-is after validation.
+    """
+    # Check if input is already a hostname (not an IP)
+    try:
+        socket.inet_aton(host)
+        is_ip = True
+    except socket.error:
+        is_ip = False
+
+    if not is_ip:
+        # Already a hostname, verify it resolves and return as-is
+        try:
+            socket.gethostbyname(host)
+            return host
+        except socket.gaierror:
+            logger.warning(f"Hostname {host} does not resolve, using anyway")
+            return host
+
+    # It's an IP address, try to resolve to hostname
+    try:
+        hostname, _, _ = socket.gethostbyaddr(host)
+        hostname = hostname.split(".")[0]
+        logger.info(f"Resolved {host} to {hostname}")
+        return hostname
+    except socket.herror:
+        pass
+
+    # Fall back to IP
+    logger.warning(f"Could not resolve {host} to hostname, using IP directly")
+    return host
+
+
+def generate_hostfile(instance: FLASHInstance, working_dir: str) -> str:
+    """Generate MPI hostfile from instance topology."""
+    hostfile_path = os.path.join(working_dir, "flash_hosts.txt")
+    with open(hostfile_path, "w") as f:
+        for _node_id, hosts in instance.hosts_by_node.items():
+            if hosts:
+                host = resolve_host(hosts[0].ip)
+                f.write(f"{host} slots={instance.ranks_per_node}\n")
+    logger.info(f"Generated hostfile at {hostfile_path}")
+    with open(hostfile_path, "r") as f:
+        logger.info(f"Hostfile contents:\n{f.read()}")
+    return hostfile_path
+
+
+def main(
+    bound_instance: BoundInstance,
+    event_sender: MpSender[Event],
+    task_receiver: MpReceiver[Task],
+) -> None:
+    """Main FLASH runner loop.
+
+    Coordinator: generates hostfile and runs mpirun (uses exo-rsh instead of SSH)
+    Workers: just report ready and wait for mpirun to spawn processes on them
+    """
+    assert isinstance(bound_instance.instance, FLASHInstance)
+    instance = bound_instance.instance
+    runner_id = bound_instance.bound_runner_id
+    my_node_id = str(bound_instance.bound_node_id)
+
+    logger.info(f"FLASH runner starting for simulation: {instance.simulation_name}")
+
+    my_rank = get_my_rank(instance, my_node_id)
+    world_size = len(instance.hosts_by_node)
+    is_coordinator = my_rank == 0
+    coordinator_ip = get_coordinator_host(instance)
+
+    logger.info(
+        f"FLASH node: rank={my_rank}, world_size={world_size}, coordinator={is_coordinator}"
+    )
+    logger.info(f"FLASH coordinator IP: {coordinator_ip}")
+
+    process: subprocess.Popen[bytes] | None = None
+    current_status: RunnerStatus = RunnerIdle()
+    shutdown_requested = False
+
+    event_sender.send(
+        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+    )
+
+    def monitor_output(proc: subprocess.Popen[bytes]) -> None:
+        """Monitor FLASH stdout for progress updates."""
+        if proc.stdout is None:
+            return
+        for line in iter(proc.stdout.readline, b""):
+            if shutdown_requested:
+                break
+            try:
+                decoded: str = line.decode("utf-8", errors="replace").strip()
+                if decoded:
+                    logger.info(f"[FLASH] {decoded}")
+            except Exception as e:
+                logger.warning(f"Error parsing FLASH output: {e}")
+
+    with task_receiver as tasks:
+        for task in tasks:
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
+            )
+            event_sender.send(TaskAcknowledged(task_id=task.task_id))
+
+            match task:
+                case LoadModel() if isinstance(current_status, RunnerIdle):
+                    current_status = RunnerLoading()
+                    logger.info("Starting FLASH simulation")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+
+                    try:
+                        if is_coordinator:
+                            # Coordinator: generate hostfile and run mpirun
+                            hostfile = generate_hostfile(
+                                instance, instance.working_directory
+                            )
+
+                            iface = instance.network_interface
+                            cmd = [
+                                MPIRUN_PATH,
+                                "-np",
+                                str(instance.total_ranks),
+                                "--hostfile",
+                                hostfile,
+                                "--wdir",
+                                instance.working_directory,
+                                "--oversubscribe",
+                                "--mca",
+                                "btl",
+                                "tcp,self",
+                                "--mca",
+                                "btl_tcp_if_include",
+                                iface,
+                                "--mca",
+                                "oob_tcp_if_include",
+                                iface,
+                                "--mca",
+                                "plm_rsh_no_tree_spawn",
+                                "1",
+                            ]
+
+                            # Use exo-rsh for remote execution (no SSH needed)
+                            cmd.extend(["--mca", "plm_rsh_agent", EXO_RSH_PATH])
+
+                            cmd.append(instance.flash_executable_path)
+
+                            logger.info(f"FLASH distributed launch: {' '.join(cmd)}")
+
+                            process = subprocess.Popen(
+                                cmd,
+                                cwd=instance.working_directory,
+                                stdout=subprocess.PIPE,
+                                stderr=subprocess.STDOUT,
+                            )
+
+                            monitor_thread = threading.Thread(
+                                target=monitor_output, args=(process,), daemon=True
+                            )
+                            monitor_thread.start()
+
+                            current_status = RunnerRunning()
+                            logger.info(
+                                f"FLASH running on {world_size} nodes with {instance.total_ranks} ranks"
+                            )
+
+                        else:
+                            # Worker: mpirun on coordinator will use exo-rsh to spawn processes here
+                            logger.info(
+                                f"Worker {my_rank}: Ready for mpirun to spawn processes via exo-rsh"
+                            )
+                            current_status = RunnerRunning()
+
+                    except Exception as e:
+                        logger.error(f"Failed to start FLASH: {e}")
+                        import traceback
+
+                        logger.error(traceback.format_exc())
+                        current_status = RunnerFailed(error_message=str(e))
+
+                case Shutdown():
+                    shutdown_requested = True
+                    current_status = RunnerShuttingDown()
+                    logger.info("FLASH runner shutting down")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+
+                    if process and process.poll() is None:
+                        logger.info("Terminating FLASH simulation")
+                        process.terminate()
+                        try:
+                            process.wait(timeout=10)
+                        except subprocess.TimeoutExpired:
+                            logger.warning("FLASH didn't terminate, killing")
+                            process.kill()
+                            process.wait()
+
+                    current_status = RunnerShutdown()
+
+                case _:
+                    if process and process.poll() is not None:
+                        exit_code = process.returncode
+                        if exit_code == 0:
+                            logger.info("FLASH simulation completed successfully")
+                            current_status = RunnerReady()
+                        else:
+                            logger.error(
+                                f"FLASH simulation failed with code {exit_code}"
+                            )
+                            current_status = RunnerFailed(
+                                error_message=f"Exit code {exit_code}"
+                            )
+
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
+            )
+            event_sender.send(
+                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+            )
+
+            if isinstance(current_status, RunnerShutdown):
+                break
+
+    if process and process.poll() is None:
+        process.terminate()
+        process.wait(timeout=5)
+
+    logger.info("FLASH runner exiting")
--- a/src/exo/plugins/registry.py
+++ b/src/exo/plugins/registry.py
@@ -0,0 +1,110 @@
+"""Plugin registry for discovering and managing plugins."""
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from loguru import logger
+
+from exo.plugins.base import ExoPlugin
+
+
+class PluginRegistry:
+    """Central registry for all plugins."""
+
+    _instance: "PluginRegistry | None" = None
+
+    def __init__(self) -> None:
+        self._plugins: dict[str, ExoPlugin] = {}
+        self._command_handlers: dict[type, ExoPlugin] = {}
+        self._instance_handlers: dict[type, ExoPlugin] = {}
+
+    @classmethod
+    def get(cls) -> "PluginRegistry":
+        """Get the singleton registry instance."""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset the singleton instance (useful for testing)."""
+        cls._instance = None
+
+    def register(self, plugin: ExoPlugin) -> None:
+        """Register a plugin and its types."""
+        if plugin.name in self._plugins:
+            raise ValueError(f"Plugin '{plugin.name}' already registered")
+
+        logger.info(f"Registering plugin: {plugin.name} v{plugin.version}")
+
+        self._plugins[plugin.name] = plugin
+
+        # Register command handlers
+        for cmd_type in plugin.get_command_types():
+            self._command_handlers[cmd_type] = plugin
+            logger.debug(f"  Registered command: {cmd_type.__name__}")
+
+        # Register instance handler
+        instance_type = plugin.get_instance_type()
+        self._instance_handlers[instance_type] = plugin
+        logger.debug(f"  Registered instance: {instance_type.__name__}")
+
+    def get_plugin(self, name: str) -> ExoPlugin | None:
+        """Get a plugin by name."""
+        return self._plugins.get(name)
+
+    def get_plugin_for_command(self, command: object) -> ExoPlugin | None:
+        """Get the plugin that handles a command."""
+        for plugin in self._plugins.values():
+            if plugin.handles_command(command):
+                return plugin
+        return None
+
+    def get_plugin_for_instance(self, instance: object) -> ExoPlugin | None:
+        """Get the plugin that manages an instance."""
+        for plugin in self._plugins.values():
+            if plugin.handles_instance(instance):
+                return plugin
+        return None
+
+    def all_plugins(self) -> Sequence[ExoPlugin]:
+        """Get all registered plugins."""
+        return list(self._plugins.values())
+
+    def get_all_api_routes(
+        self,
+    ) -> Sequence[tuple[str, str, Callable[..., Any], ExoPlugin]]:
+        """Get all API routes from all plugins."""
+        routes: list[tuple[str, str, Callable[..., Any], ExoPlugin]] = []
+        for plugin in self._plugins.values():
+            for method, path, handler in plugin.get_api_routes():
+                routes.append((method, path, handler, plugin))
+        return routes
+
+
+def discover_plugins() -> None:
+    """Auto-discover and register plugins from the implementations directory.
+
+    Plugins should have a register() function that returns an ExoPlugin instance.
+    """
+    import importlib
+    import pkgutil
+
+    registry = PluginRegistry.get()
+
+    try:
+        import exo.plugins.implementations as impl_package
+
+        for _, module_name, _ in pkgutil.iter_modules(impl_package.__path__):
+            try:
+                module = importlib.import_module(
+                    f"exo.plugins.implementations.{module_name}"
+                )
+                if hasattr(module, "register"):
+                    plugin = module.register()  # pyright: ignore[reportAny]
+                    if plugin is not None:
+                        registry.register(plugin)  # pyright: ignore[reportAny]
+            except Exception as e:
+                logger.warning(f"Failed to load plugin {module_name}: {e}")
+    except ImportError:
+        logger.debug("No plugin implementations package found")
--- a/src/exo/rsh/init.py
+++ b/src/exo/rsh/init.py
@@ -0,0 +1,13 @@
+"""Exo RSH - Remote Shell for MPI without SSH.
+
+This module provides a remote execution mechanism that allows mpirun to spawn
+processes on remote nodes without requiring SSH setup. It works by:
+
+1. Each Exo node runs an API server on port 52415 with an /execute endpoint
+2. The exo-rsh script acts as a drop-in replacement for ssh
+3. When mpirun calls "exo-rsh hostname command", it HTTP POSTs to the target's /execute
+4. The target executes the command and returns output
+
+Usage:
+    mpirun --mca plm_rsh_agent exo-rsh -np 4 --hostfile hosts.txt ./program
+"""
--- a/src/exo/rsh/client.py
+++ b/src/exo/rsh/client.py
@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""exo-rsh - Remote shell client for MPI.
+
+This script is called by mpirun as a replacement for ssh.
+Usage: exo-rsh [ssh-options...] hostname command [args...]
+
+It connects to the target node's Exo API (port 52415) and executes the command.
+"""
+
+import json
+import socket
+import sys
+from typing import Any, cast
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+# Use the same port as Exo's API server
+EXO_API_PORT = 52415
+
+
+def resolve_hostname(hostname: str) -> str:
+    """Resolve hostname to IP address."""
+    try:
+        return socket.gethostbyname(hostname)
+    except socket.gaierror:
+        # If resolution fails, try using the hostname directly
+        return hostname
+
+
+def main():
+    # Parse arguments - mpirun calls us like: exo-rsh [options] hostname command [args...]
+    # SSH options we might see: -x (disable X11), -o options, etc.
+    args = sys.argv[1:]
+
+    # Skip SSH-style options
+    hostname = None
+    command_start = 0
+
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg.startswith("-"):
+            # Skip option and its value if needed
+            if arg in ("-o", "-i", "-l", "-p", "-F"):
+                i += 2  # Skip option and its argument
+                continue
+            i += 1
+            continue
+        else:
+            # First non-option is the hostname
+            hostname = arg
+            command_start = i + 1
+            break
+        i += 1
+
+    if hostname is None or command_start >= len(args):
+        print("Usage: exo-rsh [options] hostname command [args...]", file=sys.stderr)
+        sys.exit(1)
+
+    command = args[command_start:]
+
+    # Resolve hostname to IP
+    ip = resolve_hostname(hostname)
+
+    # Make request to Exo API
+    url = f"http://{ip}:{EXO_API_PORT}/execute"
+    data = json.dumps({"command": command}).encode("utf-8")
+
+    try:
+        req = Request(url, data=data, headers={"Content-Type": "application/json"})
+        with urlopen(req, timeout=300) as response:  # pyright: ignore[reportAny]
+            response_body: bytes = cast(bytes, response.read())  # pyright: ignore[reportAny]
+            result: dict[str, Any] = json.loads(response_body.decode("utf-8"))  # pyright: ignore[reportAny]
+
+        # Output stdout/stderr
+        stdout: str = cast(str, result.get("stdout", ""))
+        stderr: str = cast(str, result.get("stderr", ""))
+        exit_code: int = cast(int, result.get("exit_code", 0))
+
+        if stdout:
+            sys.stdout.write(stdout)
+            sys.stdout.flush()
+        if stderr:
+            sys.stderr.write(stderr)
+            sys.stderr.flush()
+
+        sys.exit(exit_code)
+
+    except URLError as e:
+        print(
+            f"exo-rsh: Failed to connect to {hostname}:{EXO_API_PORT}: {e}",
+            file=sys.stderr,
+        )
+        sys.exit(255)
+    except Exception as e:
+        print(f"exo-rsh: Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -11,10 +11,8 @@ from exo.shared.types.events import (
    IndexedEvent,
    InstanceCreated,
    InstanceDeleted,
-    NodeCreated,
    NodeDownloadProgress,
-    NodeMemoryMeasured,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
    NodeTimedOut,
    RunnerDeleted,
    RunnerStatusUpdated,
@@ -27,13 +25,27 @@ from exo.shared.types.events import (
    TopologyEdgeCreated,
    TopologyEdgeDeleted,
 )
-from exo.shared.types.profiling import NodePerformanceProfile, SystemPerformanceProfile
+from exo.shared.types.profiling import (
+    NodeIdentity,
+    NodeNetworkInfo,
+    NodeThunderboltInfo,
+)
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.topology import NodeInfo
+from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
+from exo.utils.info_gatherer.info_gatherer import (
+    MacmonMetrics,
+    MacThunderboltConnections,
+    MacThunderboltIdentifiers,
+    MemoryUsage,
+    MiscData,
+    NodeConfig,
+    NodeNetworkInterfaces,
+    StaticNodeInformation,
+)


 def event_apply(event: Event, state: State) -> State:
@@ -47,16 +59,12 @@ def event_apply(event: Event, state: State) -> State:
            return apply_instance_created(event, state)
        case InstanceDeleted():
            return apply_instance_deleted(event, state)
-        case NodeCreated():
-            return apply_topology_node_created(event, state)
        case NodeTimedOut():
            return apply_node_timed_out(event, state)
-        case NodePerformanceMeasured():
-            return apply_node_performance_measured(event, state)
        case NodeDownloadProgress():
            return apply_node_download_progress(event, state)
-        case NodeMemoryMeasured():
-            return apply_node_memory_measured(event, state)
+        case NodeGatheredInfo():
+            return apply_node_gathered_info(event, state)
        case RunnerDeleted():
            return apply_runner_deleted(event, state)
        case RunnerStatusUpdated():
@@ -188,120 +196,133 @@ def apply_runner_deleted(event: RunnerDeleted, state: State) -> State:


 def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
-    topology = copy.copy(state.topology)
-    state.topology.remove_node(event.node_id)
-    node_profiles = {
-        key: value for key, value in state.node_profiles.items() if key != event.node_id
-    }
+    topology = copy.deepcopy(state.topology)
+    topology.remove_node(event.node_id)
    last_seen = {
        key: value for key, value in state.last_seen.items() if key != event.node_id
    }
+    downloads = {
+        key: value for key, value in state.downloads.items() if key != event.node_id
+    }
+    # Clean up all granular node mappings
+    node_identities = {
+        key: value
+        for key, value in state.node_identities.items()
+        if key != event.node_id
+    }
+    node_memory = {
+        key: value for key, value in state.node_memory.items() if key != event.node_id
+    }
+    node_system = {
+        key: value for key, value in state.node_system.items() if key != event.node_id
+    }
+    node_network = {
+        key: value for key, value in state.node_network.items() if key != event.node_id
+    }
+    node_thunderbolt = {
+        key: value
+        for key, value in state.node_thunderbolt.items()
+        if key != event.node_id
+    }
    return state.model_copy(
        update={
+            "downloads": downloads,
            "topology": topology,
-            "node_profiles": node_profiles,
            "last_seen": last_seen,
+            "node_identities": node_identities,
+            "node_memory": node_memory,
+            "node_system": node_system,
+            "node_network": node_network,
+            "node_thunderbolt": node_thunderbolt,
        }
    )


-def apply_node_performance_measured(
-    event: NodePerformanceMeasured, state: State
-) -> State:
-    new_profiles: Mapping[NodeId, NodePerformanceProfile] = {
-        **state.node_profiles,
-        event.node_id: event.node_profile,
-    }
-    last_seen: Mapping[NodeId, datetime] = {
-        **state.last_seen,
-        event.node_id: datetime.fromisoformat(event.when),
-    }
-    state = state.model_copy(update={"node_profiles": new_profiles})
-    topology = copy.copy(state.topology)
-    # TODO: NodeCreated
-    if not topology.contains_node(event.node_id):
-        topology.add_node(NodeInfo(node_id=event.node_id))
-    topology.update_node_profile(event.node_id, event.node_profile)
-    return state.model_copy(
-        update={
-            "node_profiles": new_profiles,
-            "topology": topology,
-            "last_seen": last_seen,
-        }
-    )
+def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
+    topology = copy.deepcopy(state.topology)
+    topology.add_node(event.node_id)
+    info = event.info

-
-def apply_node_memory_measured(event: NodeMemoryMeasured, state: State) -> State:
-    existing = state.node_profiles.get(event.node_id)
-    topology = copy.copy(state.topology)
-
-    if existing is None:
-        created = NodePerformanceProfile(
-            model_id="unknown",
-            chip_id="unknown",
-            friendly_name="Unknown",
-            memory=event.memory,
-            network_interfaces=[],
-            system=SystemPerformanceProfile(
-                # TODO: flops_fp16=0.0,
-                gpu_usage=0.0,
-                temp=0.0,
-                sys_power=0.0,
-                pcpu_usage=0.0,
-                ecpu_usage=0.0,
-                ane_power=0.0,
-            ),
-        )
-        created_profiles: Mapping[NodeId, NodePerformanceProfile] = {
-            **state.node_profiles,
-            event.node_id: created,
-        }
-        last_seen: Mapping[NodeId, datetime] = {
+    # Build update dict with only the mappings that change
+    update: dict[str, object] = {
+        "last_seen": {
            **state.last_seen,
            event.node_id: datetime.fromisoformat(event.when),
-        }
-        if not topology.contains_node(event.node_id):
-            topology.add_node(NodeInfo(node_id=event.node_id))
-            # TODO: NodeCreated
-        topology.update_node_profile(event.node_id, created)
-        return state.model_copy(
-            update={
-                "node_profiles": created_profiles,
-                "topology": topology,
-                "last_seen": last_seen,
-            }
-        )
-
-    updated = existing.model_copy(update={"memory": event.memory})
-    updated_profiles: Mapping[NodeId, NodePerformanceProfile] = {
-        **state.node_profiles,
-        event.node_id: updated,
+        },
+        "topology": topology,
    }
-    # TODO: NodeCreated
-    if not topology.contains_node(event.node_id):
-        topology.add_node(NodeInfo(node_id=event.node_id))
-    topology.update_node_profile(event.node_id, updated)
-    return state.model_copy(
-        update={"node_profiles": updated_profiles, "topology": topology}
-    )

+    match info:
+        case MacmonMetrics():
+            update["node_system"] = {
+                **state.node_system,
+                event.node_id: info.system_profile,
+            }
+            update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
+        case MemoryUsage():
+            update["node_memory"] = {**state.node_memory, event.node_id: info}
+        case NodeConfig():
+            pass
+        case MiscData():
+            current_identity = state.node_identities.get(event.node_id, NodeIdentity())
+            new_identity = current_identity.model_copy(
+                update={"friendly_name": info.friendly_name}
+            )
+            update["node_identities"] = {
+                **state.node_identities,
+                event.node_id: new_identity,
+            }
+        case StaticNodeInformation():
+            current_identity = state.node_identities.get(event.node_id, NodeIdentity())
+            new_identity = current_identity.model_copy(
+                update={"model_id": info.model, "chip_id": info.chip}
+            )
+            update["node_identities"] = {
+                **state.node_identities,
+                event.node_id: new_identity,
+            }
+        case NodeNetworkInterfaces():
+            update["node_network"] = {
+                **state.node_network,
+                event.node_id: NodeNetworkInfo(interfaces=info.ifaces),
+            }
+        case MacThunderboltIdentifiers():
+            update["node_thunderbolt"] = {
+                **state.node_thunderbolt,
+                event.node_id: NodeThunderboltInfo(interfaces=info.idents),
+            }
+        case MacThunderboltConnections():
+            conn_map = {
+                tb_ident.domain_uuid: (nid, tb_ident.rdma_interface)
+                for nid in state.node_thunderbolt
+                for tb_ident in state.node_thunderbolt[nid].interfaces
+            }
+            as_rdma_conns = [
+                Connection(
+                    source=event.node_id,
+                    sink=conn_map[tb_conn.sink_uuid][0],
+                    edge=RDMAConnection(
+                        source_rdma_iface=conn_map[tb_conn.source_uuid][1],
+                        sink_rdma_iface=conn_map[tb_conn.sink_uuid][1],
+                    ),
+                )
+                for tb_conn in info.conns
+                if tb_conn.source_uuid in conn_map
+                if tb_conn.sink_uuid in conn_map
+            ]
+            topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)

-def apply_topology_node_created(event: NodeCreated, state: State) -> State:
-    topology = copy.copy(state.topology)
-    topology.add_node(NodeInfo(node_id=event.node_id))
-    return state.model_copy(update={"topology": topology})
+    return state.model_copy(update=update)


 def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
-    topology = copy.copy(state.topology)
-    topology.add_connection(event.edge)
+    topology = copy.deepcopy(state.topology)
+    topology.add_connection(event.conn)
    return state.model_copy(update={"topology": topology})


 def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
-    topology = copy.copy(state.topology)
-    if not topology.contains_connection(event.edge):
-        return state
-    topology.remove_connection(event.edge)
+    topology = copy.deepcopy(state.topology)
+    topology.remove_connection(event.conn)
    # TODO: Clean up removing the reverse connection
    return state.model_copy(update={"topology": topology})
--- a/src/exo/shared/constants.py
+++ b/src/exo/shared/constants.py
@@ -38,6 +38,7 @@ EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"

 # Identity (config)
 EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"
+EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml"

 # libp2p topics for event forwarding
 LIBP2P_LOCAL_EVENTS_TOPIC = "worker_events"
--- a/src/exo/shared/logging.py
+++ b/src/exo/shared/logging.py
@@ -11,9 +11,6 @@ class InterceptLogger(HypercornLogger):
    def __init__(self, config: Config):
        super().__init__(config)
        assert self.error_logger
-        # TODO: Decide if we want to provide access logs
-        # assert self.access_logger
-        # self.access_logger.handlers = [_InterceptHandler()]
        self.error_logger.handlers = [_InterceptHandler()]


--- a/src/exo/shared/tests/test_apply/test_apply_node_download.py
+++ b/src/exo/shared/tests/test_apply/test_apply_node_download.py
@@ -43,7 +43,4 @@ def test_apply_two_node_download_progress():
        NodeDownloadProgress(download_progress=event2), state
    )

-    # TODO: This test is failing. We should support the following:
-    # 1. Downloading multiple models concurrently on the same node (one per runner is fine).
-    # 2. Downloading a model, it completes, then downloading a different model on the same node.
    assert new_state.downloads == {NodeId("node-1"): [event1, event2]}
--- a/src/exo/shared/tests/test_state_serialization.py
+++ b/src/exo/shared/tests/test_state_serialization.py
@@ -1,7 +1,7 @@
 from exo.shared.types.common import NodeId
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
-from exo.shared.types.topology import Connection
+from exo.shared.types.topology import Connection, SocketConnection


 def test_state_serialization_roundtrip() -> None:
@@ -12,9 +12,11 @@ def test_state_serialization_roundtrip() -> None:
    node_b = NodeId("node-b")

    connection = Connection(
-        local_node_id=node_a,
-        send_back_node_id=node_b,
-        send_back_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
+        source=node_a,
+        sink=node_b,
+        edge=SocketConnection(
+            sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
+        ),
    )

    state = State()
@@ -23,5 +25,11 @@ def test_state_serialization_roundtrip() -> None:
    json_repr = state.model_dump_json()
    restored_state = State.model_validate_json(json_repr)

-    assert state.topology.to_snapshot() == restored_state.topology.to_snapshot()
+    assert (
+        state.topology.to_snapshot().nodes
+        == restored_state.topology.to_snapshot().nodes
+    )
+    assert set(state.topology.to_snapshot().connections) == set(
+        restored_state.topology.to_snapshot().connections
+    )
    assert restored_state.model_dump_json() == json_repr
--- a/src/exo/shared/topology.py
+++ b/src/exo/shared/topology.py
@@ -1,203 +1,227 @@
 import contextlib
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
 from typing import Iterable

 import rustworkx as rx
 from pydantic import BaseModel, ConfigDict

 from exo.shared.types.common import NodeId
-from exo.shared.types.profiling import ConnectionProfile, NodePerformanceProfile
-from exo.shared.types.topology import Connection, NodeInfo
+from exo.shared.types.topology import (
+    Connection,
+    Cycle,
+    RDMAConnection,
+    SocketConnection,
+)


 class TopologySnapshot(BaseModel):
-    nodes: list[NodeInfo]
-    connections: list[Connection]
+    nodes: Sequence[NodeId]
+    connections: Mapping[
+        NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]
+    ]

-    model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")


+@dataclass
 class Topology:
-    def __init__(self) -> None:
-        self._graph: rx.PyDiGraph[NodeInfo, Connection] = rx.PyDiGraph()
-        self._node_id_to_rx_id_map: dict[NodeId, int] = dict()
-        self._rx_id_to_node_id_map: dict[int, NodeId] = dict()
-        self._edge_id_to_rx_id_map: dict[Connection, int] = dict()
+    _graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = field(
+        init=False, default_factory=rx.PyDiGraph
+    )
+    _vertex_indices: dict[NodeId, int] = field(init=False, default_factory=dict)

    def to_snapshot(self) -> TopologySnapshot:
        return TopologySnapshot(
-            nodes=list(self.list_nodes()),
-            connections=list(self.list_connections()),
+            nodes=list(self.list_nodes()), connections=self.map_connections()
        )

    @classmethod
    def from_snapshot(cls, snapshot: TopologySnapshot) -> "Topology":
        topology = cls()

-        for node in snapshot.nodes:
+        for node_id in snapshot.nodes:
            with contextlib.suppress(ValueError):
-                topology.add_node(node)
+                topology.add_node(node_id)

-        for connection in snapshot.connections:
-            topology.add_connection(connection)
+        for source in snapshot.connections:
+            for sink in snapshot.connections[source]:
+                for edge in snapshot.connections[source][sink]:
+                    topology.add_connection(
+                        Connection(source=source, sink=sink, edge=edge)
+                    )

        return topology

-    def add_node(self, node: NodeInfo) -> None:
-        if node.node_id in self._node_id_to_rx_id_map:
+    def add_node(self, node_id: NodeId) -> None:
+        if node_id in self._vertex_indices:
            return
-        rx_id = self._graph.add_node(node)
-        self._node_id_to_rx_id_map[node.node_id] = rx_id
-        self._rx_id_to_node_id_map[rx_id] = node.node_id
+        rx_id = self._graph.add_node(node_id)
+        self._vertex_indices[node_id] = rx_id

    def node_is_leaf(self, node_id: NodeId) -> bool:
        return (
-            node_id in self._node_id_to_rx_id_map
-            and len(self._graph.neighbors(self._node_id_to_rx_id_map[node_id])) == 1
+            node_id in self._vertex_indices
+            and len(self._graph.neighbors(self._vertex_indices[node_id])) <= 1
        )

    def neighbours(self, node_id: NodeId) -> list[NodeId]:
        return [
-            self._rx_id_to_node_id_map[rx_id]
-            for rx_id in self._graph.neighbors(self._node_id_to_rx_id_map[node_id])
+            self._graph[rx_id]
+            for rx_id in self._graph.neighbors(self._vertex_indices[node_id])
        ]

-    def out_edges(self, node_id: NodeId) -> list[tuple[NodeId, Connection]]:
-        if node_id not in self._node_id_to_rx_id_map:
+    def out_edges(self, node_id: NodeId) -> Iterable[Connection]:
+        if node_id not in self._vertex_indices:
            return []
-        return [
-            (self._rx_id_to_node_id_map[nid], conn)
-            for _, nid, conn in self._graph.out_edges(
-                self._node_id_to_rx_id_map[node_id]
+        return (
+            Connection(source=self._graph[source], sink=self._graph[sink], edge=edge)
+            for source, sink, edge in self._graph.out_edges(
+                self._vertex_indices[node_id]
            )
-        ]
+        )

    def contains_node(self, node_id: NodeId) -> bool:
-        return node_id in self._node_id_to_rx_id_map
+        return node_id in self._vertex_indices

-    def contains_connection(self, connection: Connection) -> bool:
-        return connection in self._edge_id_to_rx_id_map
-
-    def add_connection(
-        self,
-        connection: Connection,
-    ) -> None:
-        if connection.local_node_id not in self._node_id_to_rx_id_map:
-            self.add_node(NodeInfo(node_id=connection.local_node_id))
-        if connection.send_back_node_id not in self._node_id_to_rx_id_map:
-            self.add_node(NodeInfo(node_id=connection.send_back_node_id))
-
-        if connection in self._edge_id_to_rx_id_map:
+    def add_connection(self, conn: Connection) -> None:
+        source, sink, edge = conn.source, conn.sink, conn.edge
+        del conn
+        if edge in self.get_all_connections_between(source, sink):
            return

-        src_id = self._node_id_to_rx_id_map[connection.local_node_id]
-        sink_id = self._node_id_to_rx_id_map[connection.send_back_node_id]
+        if source not in self._vertex_indices:
+            self.add_node(source)
+        if sink not in self._vertex_indices:
+            self.add_node(sink)

-        rx_id = self._graph.add_edge(src_id, sink_id, connection)
-        self._edge_id_to_rx_id_map[connection] = rx_id
+        src_id = self._vertex_indices[source]
+        sink_id = self._vertex_indices[sink]

-    def list_nodes(self) -> Iterable[NodeInfo]:
-        return (self._graph[i] for i in self._graph.node_indices())
+        _ = self._graph.add_edge(src_id, sink_id, edge)

-    def list_connections(self) -> Iterable[Connection]:
-        return (connection for _, _, connection in self._graph.weighted_edge_list())
+    def get_all_connections_between(
+        self, source: NodeId, sink: NodeId
+    ) -> Iterable[SocketConnection | RDMAConnection]:
+        if source not in self._vertex_indices:
+            return []
+        if sink not in self._vertex_indices:
+            return []

-    def get_node_profile(self, node_id: NodeId) -> NodePerformanceProfile | None:
+        src_id = self._vertex_indices[source]
+        sink_id = self._vertex_indices[sink]
        try:
-            rx_idx = self._node_id_to_rx_id_map[node_id]
-            return self._graph.get_node_data(rx_idx).node_profile
-        except KeyError:
-            return None
+            return self._graph.get_all_edge_data(src_id, sink_id)
+        except rx.NoEdgeBetweenNodes:
+            return []

-    def update_node_profile(
-        self, node_id: NodeId, node_profile: NodePerformanceProfile
-    ) -> None:
-        rx_idx = self._node_id_to_rx_id_map[node_id]
-        self._graph[rx_idx].node_profile = node_profile
+    def list_nodes(self) -> Iterable[NodeId]:
+        return self._graph.nodes()

-    def update_connection_profile(self, connection: Connection) -> None:
-        rx_idx = self._edge_id_to_rx_id_map[connection]
-        self._graph.update_edge_by_index(rx_idx, connection)
+    def map_connections(
+        self,
+    ) -> Mapping[NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]]:
+        base: dict[NodeId, dict[NodeId, list[SocketConnection | RDMAConnection]]] = {}
+        for src_id, sink_id, connection in self._graph.weighted_edge_list():
+            source = self._graph[src_id]
+            sink = self._graph[sink_id]
+            if source not in base:
+                base[source] = {}
+            if sink not in base[source]:
+                base[source][sink] = []
+            base[source][sink].append(connection)
+        return base

-    def get_connection_profile(
-        self, connection: Connection
-    ) -> ConnectionProfile | None:
-        try:
-            rx_idx = self._edge_id_to_rx_id_map[connection]
-            return self._graph.get_edge_data_by_index(rx_idx).connection_profile
-        except KeyError:
-            return None
+    def list_connections(
+        self,
+    ) -> Iterable[Connection]:
+        return (
+            (
+                Connection(
+                    source=self._graph[src_id],
+                    sink=self._graph[sink_id],
+                    edge=connection,
+                )
+            )
+            for src_id, sink_id, connection in self._graph.weighted_edge_list()
+        )

    def remove_node(self, node_id: NodeId) -> None:
-        if node_id not in self._node_id_to_rx_id_map:
+        if node_id not in self._vertex_indices:
            return

-        for connection in self.list_connections():
-            if (
-                connection.local_node_id == node_id
-                or connection.send_back_node_id == node_id
-            ):
-                self.remove_connection(connection)
-
-        rx_idx = self._node_id_to_rx_id_map[node_id]
+        rx_idx = self._vertex_indices[node_id]
        self._graph.remove_node(rx_idx)

-        del self._node_id_to_rx_id_map[node_id]
-        del self._rx_id_to_node_id_map[rx_idx]
+        del self._vertex_indices[node_id]

-    def remove_connection(self, connection: Connection) -> None:
-        if connection not in self._edge_id_to_rx_id_map:
+    def replace_all_out_rdma_connections(
+        self, source: NodeId, new_connections: Sequence[Connection]
+    ) -> None:
+        for conn_idx in self._graph.out_edge_indices(self._vertex_indices[source]):
+            if isinstance(self._graph.get_edge_data_by_index(conn_idx), RDMAConnection):
+                self._graph.remove_edge_from_index(conn_idx)
+        for conn in new_connections:
+            self.add_connection(conn)
+
+    def remove_connection(self, conn: Connection) -> None:
+        if (
+            conn.source not in self._vertex_indices
+            or conn.sink not in self._vertex_indices
+        ):
            return
-        rx_idx = self._edge_id_to_rx_id_map[connection]
-        self._graph.remove_edge_from_index(rx_idx)
-        del self._edge_id_to_rx_id_map[connection]
+        for conn_idx in self._graph.edge_indices_from_endpoints(
+            self._vertex_indices[conn.source], self._vertex_indices[conn.sink]
+        ):
+            if self._graph.get_edge_data_by_index(conn_idx) == conn.edge:
+                self._graph.remove_edge_from_index(conn_idx)
+
+    def get_cycles(self) -> list[Cycle]:
+        """Get simple cycles in the graph, including singleton cycles"""

-    def get_cycles(self) -> list[list[NodeInfo]]:
        cycle_idxs = rx.simple_cycles(self._graph)
-        cycles: list[list[NodeInfo]] = []
+        cycles: list[Cycle] = []
        for cycle_idx in cycle_idxs:
-            cycle = [self._graph[idx] for idx in cycle_idx]
+            cycle = Cycle(node_ids=[self._graph[idx] for idx in cycle_idx])
            cycles.append(cycle)
-
+        for node_id in self.list_nodes():
+            cycles.append(Cycle(node_ids=[node_id]))
        return cycles

-    def get_cycles_tb(self) -> list[list[NodeInfo]]:
+    def get_cycles_tb(self) -> list[Cycle]:
        tb_edges = [
            (u, v, conn)
            for u, v, conn in self._graph.weighted_edge_list()
            if conn.is_thunderbolt()
        ]

-        tb_graph: rx.PyDiGraph[NodeInfo, Connection] = rx.PyDiGraph()
+        tb_graph: rx.PyDiGraph[NodeId, SocketConnection] = rx.PyDiGraph()
        tb_graph.add_nodes_from(self._graph.nodes())

        for u, v, conn in tb_edges:
-            tb_graph.add_edge(u, v, conn)
+            if isinstance(conn, SocketConnection):
+                tb_graph.add_edge(u, v, conn)

        cycle_idxs = rx.simple_cycles(tb_graph)
-        cycles: list[list[NodeInfo]] = []
+        cycles: list[Cycle] = []
        for cycle_idx in cycle_idxs:
-            cycle = [tb_graph[idx] for idx in cycle_idx]
+            cycle = Cycle(node_ids=[tb_graph[idx] for idx in cycle_idx])
            cycles.append(cycle)

        return cycles

-    def get_subgraph_from_nodes(self, nodes: list[NodeInfo]) -> "Topology":
-        node_idxs = [node.node_id for node in nodes]
-        rx_idxs = [self._node_id_to_rx_id_map[idx] for idx in node_idxs]
+    def get_subgraph_from_nodes(self, node_ids: list[NodeId]) -> "Topology":
        topology = Topology()
-        for rx_idx in rx_idxs:
-            topology.add_node(self._graph[rx_idx])
+        for node_id in node_ids:
+            topology.add_node(node_id)
        for connection in self.list_connections():
-            if (
-                connection.local_node_id in node_idxs
-                and connection.send_back_node_id in node_idxs
-            ):
+            if connection.source in node_ids and connection.sink in node_ids:
                topology.add_connection(connection)
        return topology

-    def is_thunderbolt_cycle(self, cycle: list[NodeInfo]) -> bool:
-        node_idxs = [node.node_id for node in cycle]
-        rx_idxs = [self._node_id_to_rx_id_map[idx] for idx in node_idxs]
+    def is_thunderbolt_cycle(self, cycle: Cycle) -> bool:
+        node_idxs = [node for node in cycle]
+        rx_idxs = [self._vertex_indices[idx] for idx in node_idxs]
        for rid in rx_idxs:
            for neighbor_rid in self._graph.neighbors(rid):
                if neighbor_rid not in rx_idxs:
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -157,13 +157,10 @@ 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):
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -1,6 +1,6 @@
 from enum import Enum

-from exo.shared.types.api import GenerationStats, TopLogprobItem
+from exo.shared.types.api import GenerationStats
 from exo.utils.pydantic_ext import TaggedModel

 from .api import FinishReason
@@ -20,8 +20,6 @@ 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
--- a/src/exo/shared/types/claude_api.py
+++ b/src/exo/shared/types/claude_api.py
@@ -1,168 +0,0 @@
-"""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
-)
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -1,8 +1,8 @@
 from pydantic import Field

+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,7 +17,7 @@ class TestCommand(BaseCommand):


 class ChatCompletion(BaseCommand):
-    request_params: ResponsesRequest
+    request_params: ChatCompletionTaskParams


 class PlaceInstance(BaseCommand):
@@ -35,6 +35,26 @@ class DeleteInstance(BaseCommand):
    instance_id: InstanceId


+class LaunchFLASH(BaseCommand):
+    """Command to launch a FLASH MPI simulation."""
+
+    simulation_name: str
+    flash_executable_path: str
+    parameter_file_path: str
+    working_directory: str
+    ranks_per_node: int = 1
+    min_nodes: int = 1
+    # Optional: explicit hostnames for MPI (e.g., "s14,james21-1")
+    # Used when topology edges don't contain IP addresses
+    hosts: str = ""
+
+
+class StopFLASH(BaseCommand):
+    """Command to stop a running FLASH simulation."""
+
+    instance_id: InstanceId
+
+
 class TaskFinished(BaseCommand):
    finished_command_id: CommandId

@@ -50,6 +70,8 @@ Command = (
    | PlaceInstance
    | CreateInstance
    | DeleteInstance
+    | LaunchFLASH
+    | StopFLASH
    | TaskFinished
 )

--- a/src/exo/shared/types/events.py
+++ b/src/exo/shared/types/events.py
@@ -2,14 +2,14 @@ from datetime import datetime

 from pydantic import Field

-from exo.shared.topology import Connection, NodePerformanceProfile
+from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
-from exo.shared.types.profiling import MemoryPerformanceProfile
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
+from exo.utils.info_gatherer.info_gatherer import GatheredInfo
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel


@@ -76,25 +76,15 @@ class RunnerDeleted(BaseEvent):
    runner_id: RunnerId


-# TODO
-class NodeCreated(BaseEvent):
-    node_id: NodeId
-
-
 class NodeTimedOut(BaseEvent):
    node_id: NodeId


-class NodePerformanceMeasured(BaseEvent):
+# TODO: bikeshed this name
+class NodeGatheredInfo(BaseEvent):
    node_id: NodeId
    when: str  # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device
-    node_profile: NodePerformanceProfile
-
-
-class NodeMemoryMeasured(BaseEvent):
-    node_id: NodeId
-    when: str  # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device
-    memory: MemoryPerformanceProfile
+    info: GatheredInfo


 class NodeDownloadProgress(BaseEvent):
@@ -107,11 +97,11 @@ class ChunkGenerated(BaseEvent):


 class TopologyEdgeCreated(BaseEvent):
-    edge: Connection
+    conn: Connection


 class TopologyEdgeDeleted(BaseEvent):
-    edge: Connection
+    conn: Connection


 Event = (
@@ -125,10 +115,8 @@ Event = (
    | InstanceDeleted
    | RunnerStatusUpdated
    | RunnerDeleted
-    | NodeCreated
    | NodeTimedOut
-    | NodePerformanceMeasured
-    | NodeMemoryMeasured
+    | NodeGatheredInfo
    | NodeDownloadProgress
    | ChunkGenerated
    | TopologyEdgeCreated
--- a/src/exo/shared/types/multiaddr.py
+++ b/src/exo/shared/types/multiaddr.py
@@ -1,10 +1,11 @@
 import re
 from typing import ClassVar

-from pydantic import BaseModel, computed_field, field_validator
+from pydantic import BaseModel, ConfigDict, computed_field, field_validator


 class Multiaddr(BaseModel):
+    model_config = ConfigDict(frozen=True)
    address: str

    PATTERNS: ClassVar[list[str]] = [
--- a/src/exo/shared/types/openai_responses.py
+++ b/src/exo/shared/types/openai_responses.py
@@ -1,190 +0,0 @@
-"""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
-)
--- a/src/exo/shared/types/profiling.py
+++ b/src/exo/shared/types/profiling.py
@@ -1,12 +1,14 @@
+from collections.abc import Sequence
 from typing import Self

 import psutil

 from exo.shared.types.memory import Memory
+from exo.shared.types.thunderbolt import ThunderboltIdentifier
 from exo.utils.pydantic_ext import CamelCaseModel


-class MemoryPerformanceProfile(CamelCaseModel):
+class MemoryUsage(CamelCaseModel):
    ram_total: Memory
    ram_available: Memory
    swap_total: Memory
@@ -44,7 +46,6 @@ class SystemPerformanceProfile(CamelCaseModel):
    sys_power: float = 0.0
    pcpu_usage: float = 0.0
    ecpu_usage: float = 0.0
-    ane_power: float = 0.0


 class NetworkInterfaceInfo(CamelCaseModel):
@@ -52,16 +53,21 @@ class NetworkInterfaceInfo(CamelCaseModel):
    ip_address: str


-class NodePerformanceProfile(CamelCaseModel):
-    model_id: str
-    chip_id: str
-    friendly_name: str
-    memory: MemoryPerformanceProfile
-    network_interfaces: list[NetworkInterfaceInfo] = []
-    system: SystemPerformanceProfile
+class NodeIdentity(CamelCaseModel):
+    """Static and slow-changing node identification data."""
+
+    model_id: str = "Unknown"
+    chip_id: str = "Unknown"
+    friendly_name: str = "Unknown"


-class ConnectionProfile(CamelCaseModel):
-    throughput: float
-    latency: float
-    jitter: float
+class NodeNetworkInfo(CamelCaseModel):
+    """Network interface information for a node."""
+
+    interfaces: Sequence[NetworkInterfaceInfo] = []
+
+
+class NodeThunderboltInfo(CamelCaseModel):
+    """Thunderbolt interface identifiers for a node."""
+
+    interfaces: Sequence[ThunderboltIdentifier] = []
--- a/src/exo/shared/types/state.py
+++ b/src/exo/shared/types/state.py
@@ -7,7 +7,13 @@ from pydantic.alias_generators import to_camel

 from exo.shared.topology import Topology, TopologySnapshot
 from exo.shared.types.common import NodeId
-from exo.shared.types.profiling import NodePerformanceProfile
+from exo.shared.types.profiling import (
+    MemoryUsage,
+    NodeIdentity,
+    NodeNetworkInfo,
+    NodeThunderboltInfo,
+    SystemPerformanceProfile,
+)
 from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -35,11 +41,17 @@ class State(CamelCaseModel):
    runners: Mapping[RunnerId, RunnerStatus] = {}
    downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
    tasks: Mapping[TaskId, Task] = {}
-    node_profiles: Mapping[NodeId, NodePerformanceProfile] = {}
    last_seen: Mapping[NodeId, datetime] = {}
    topology: Topology = Field(default_factory=Topology)
    last_event_applied_idx: int = Field(default=-1, ge=-1)

+    # Granular node state mappings (update independently at different frequencies)
+    node_identities: Mapping[NodeId, NodeIdentity] = {}
+    node_memory: Mapping[NodeId, MemoryUsage] = {}
+    node_system: Mapping[NodeId, SystemPerformanceProfile] = {}
+    node_network: Mapping[NodeId, NodeNetworkInfo] = {}
+    node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
+
    @field_serializer("topology", mode="plain")
    def _encode_topology(self, value: Topology) -> TopologySnapshot:
        return value.to_snapshot()
--- a/src/exo/shared/types/tasks.py
+++ b/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: ResponsesRequest
+    task_params: ChatCompletionTaskParams

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


-class NodeInfo(CamelCaseModel):
-    node_id: NodeId
-    node_profile: NodePerformanceProfile | None = None
+@dataclass(frozen=True)
+class Cycle:
+    node_ids: list[NodeId]
+
+    def __len__(self) -> int:
+        return self.node_ids.__len__()
+
+    def __iter__(self) -> Iterator[NodeId]:
+        return self.node_ids.__iter__()


-class Connection(CamelCaseModel):
-    local_node_id: NodeId
-    send_back_node_id: NodeId
-    send_back_multiaddr: Multiaddr
-    connection_profile: ConnectionProfile | None = None
-
-    def __hash__(self) -> int:
-        return hash(
-            (
-                self.local_node_id,
-                self.send_back_node_id,
-                self.send_back_multiaddr.address,
-            )
-        )
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Connection):
-            raise ValueError("Cannot compare Connection with non-Connection")
-        return (
-            self.local_node_id == other.local_node_id
-            and self.send_back_node_id == other.send_back_node_id
-            and self.send_back_multiaddr == other.send_back_multiaddr
-        )
+class RDMAConnection(FrozenModel):
+    source_rdma_iface: str
+    sink_rdma_iface: str

    def is_thunderbolt(self) -> bool:
-        return str(self.send_back_multiaddr.ipv4_address).startswith("169.254")
+        return True
+
+
+class SocketConnection(FrozenModel):
+    sink_multiaddr: Multiaddr
+
+    def __hash__(self):
+        return hash(self.sink_multiaddr.ip_address)
+
+    def is_thunderbolt(self) -> bool:
+        return str(self.sink_multiaddr.ipv4_address).startswith("169.254")
+
+
+class Connection(FrozenModel):
+    source: NodeId
+    sink: NodeId
+    edge: RDMAConnection | SocketConnection
--- a/src/exo/shared/types/worker/instances.py
+++ b/src/exo/shared/types/worker/instances.py
@@ -14,6 +14,7 @@ class InstanceId(Id):
 class InstanceMeta(str, Enum):
    MlxRing = "MlxRing"
    MlxJaccl = "MlxJaccl"
+    FLASH = "FLASH"


 class BaseInstance(TaggedModel):
@@ -30,12 +31,31 @@ class MlxRingInstance(BaseInstance):


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


+class FLASHInstance(BaseInstance):
+    """Instance for FLASH MPI simulation.
+
+    Unlike MLX instances which do tensor parallelism, FLASH instances
+    coordinate MPI processes across nodes. Each node runs one or more
+    MPI ranks of the FLASH simulation.
+    """
+
+    hosts_by_node: dict[NodeId, list[Host]]
+    flash_executable_path: str
+    parameter_file_path: str
+    working_directory: str
+    ranks_per_node: int = 1
+    total_ranks: int
+    simulation_name: str
+    coordinator_ip: str
+    network_interface: str = "en0"  # Network interface for MPI (e.g., en0, eth0)
+
+
 # TODO: Single node instance
-Instance = MlxRingInstance | MlxJacclInstance
+Instance = MlxRingInstance | MlxJacclInstance | FLASHInstance


 class BoundInstance(CamelCaseModel):
--- a/src/exo/shared/types/worker/resource_monitor.py
+++ b/src/exo/shared/types/worker/resource_monitor.py
@@ -1,43 +0,0 @@
-import asyncio
-from abc import ABC, abstractmethod
-from collections.abc import Coroutine
-from typing import Callable
-
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    SystemPerformanceProfile,
-)
-
-
-class ResourceCollector(ABC):
-    @abstractmethod
-    async def collect(self) -> SystemPerformanceProfile | MemoryPerformanceProfile: ...
-
-
-class SystemResourceCollector(ResourceCollector):
-    async def collect(self) -> SystemPerformanceProfile: ...
-
-
-class MemoryResourceCollector(ResourceCollector):
-    async def collect(self) -> MemoryPerformanceProfile: ...
-
-
-class ResourceMonitor:
-    data_collectors: list[ResourceCollector]
-    effect_handlers: set[
-        Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]
-    ]
-
-    async def _collect(
-        self,
-    ) -> list[SystemPerformanceProfile | MemoryPerformanceProfile]:
-        tasks: list[
-            Coroutine[None, None, SystemPerformanceProfile | MemoryPerformanceProfile]
-        ] = [collector.collect() for collector in self.data_collectors]
-        return await asyncio.gather(*tasks)
-
-    async def collect(self) -> None:
-        profiles = await self._collect()
-        for profile in profiles:
-            for effect_handler in self.effect_handlers:
-                effect_handler(profile)
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/src/exo/shared/types/worker/runner_response.py
@@ -1,4 +1,4 @@
-from exo.shared.types.api import FinishReason, GenerationStats, TopLogprobItem
+from exo.shared.types.api import FinishReason, GenerationStats
 from exo.utils.pydantic_ext import TaggedModel


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

--- a/src/exo/utils/info_gatherer/init.py
+++ b/src/exo/utils/info_gatherer/init.py
--- a/src/exo/utils/info_gatherer/info_gatherer.py
+++ b/src/exo/utils/info_gatherer/info_gatherer.py
@@ -0,0 +1,235 @@
+import os
+import shutil
+import sys
+import tomllib
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from subprocess import CalledProcessError
+from typing import Self, cast
+
+import anyio
+from anyio import create_task_group, open_process
+from anyio.abc import TaskGroup
+from anyio.streams.buffered import BufferedByteReceiveStream
+from anyio.streams.text import TextReceiveStream
+from loguru import logger
+
+from exo.shared.constants import EXO_CONFIG_FILE
+from exo.shared.types.memory import Memory
+from exo.shared.types.profiling import (
+    MemoryUsage,
+    NetworkInterfaceInfo,
+)
+from exo.shared.types.thunderbolt import (
+    ThunderboltConnection,
+    ThunderboltConnectivity,
+    ThunderboltIdentifier,
+)
+from exo.utils.channels import Sender
+from exo.utils.pydantic_ext import TaggedModel
+
+from .macmon import MacmonMetrics
+from .system_info import get_friendly_name, get_model_and_chip, get_network_interfaces
+
+IS_DARWIN = sys.platform == "darwin"
+
+
+class StaticNodeInformation(TaggedModel):
+    """Node information that should NEVER change, to be gathered once at startup"""
+
+    model: str
+    chip: str
+
+    @classmethod
+    async def gather(cls) -> Self:
+        model, chip = await get_model_and_chip()
+        return cls(model=model, chip=chip)
+
+
+class NodeNetworkInterfaces(TaggedModel):
+    ifaces: Sequence[NetworkInterfaceInfo]
+
+
+class MacThunderboltIdentifiers(TaggedModel):
+    idents: Sequence[ThunderboltIdentifier]
+
+
+class MacThunderboltConnections(TaggedModel):
+    conns: Sequence[ThunderboltConnection]
+
+
+class NodeConfig(TaggedModel):
+    """Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there"""
+
+    @classmethod
+    async def gather(cls) -> Self | None:
+        cfg_file = anyio.Path(EXO_CONFIG_FILE)
+        await cfg_file.touch(exist_ok=True)
+        async with await cfg_file.open("rb") as f:
+            try:
+                contents = (await f.read()).decode("utf-8")
+                data = tomllib.loads(contents)
+                return cls.model_validate(data)
+            except (tomllib.TOMLDecodeError, UnicodeDecodeError):
+                logger.warning("Invalid config file, skipping...")
+                return None
+
+
+class MiscData(TaggedModel):
+    """Node information that may slowly change that doesn't fall into the other categories"""
+
+    friendly_name: str
+
+    @classmethod
+    async def gather(cls) -> Self:
+        return cls(friendly_name=await get_friendly_name())
+
+
+async def _gather_iface_map() -> dict[str, str] | None:
+    proc = await anyio.run_process(
+        ["networksetup", "-listallhardwareports"], check=False
+    )
+    if proc.returncode != 0:
+        return None
+
+    ports: dict[str, str] = {}
+    port = ""
+    for line in proc.stdout.decode("utf-8").split("\n"):
+        if line.startswith("Hardware Port:"):
+            port = line.split(": ")[1]
+        elif line.startswith("Device:"):
+            ports[port] = line.split(": ")[1]
+            port = ""
+    if "" in ports:
+        del ports[""]
+    return ports
+
+
+GatheredInfo = (
+    MacmonMetrics
+    | MemoryUsage
+    | NodeNetworkInterfaces
+    | MacThunderboltIdentifiers
+    | MacThunderboltConnections
+    | NodeConfig
+    | MiscData
+    | StaticNodeInformation
+)
+
+
+@dataclass
+class InfoGatherer:
+    info_sender: Sender[GatheredInfo]
+    interface_watcher_interval: float | None = 10
+    misc_poll_interval: float | None = 60
+    system_profiler_interval: float | None = 5 if IS_DARWIN else None
+    memory_poll_rate: float | None = None if IS_DARWIN else 1
+    macmon_interval: float | None = 1 if IS_DARWIN else None
+    _tg: TaskGroup = field(init=False, default_factory=create_task_group)
+
+    async def run(self):
+        async with self._tg as tg:
+            if IS_DARWIN:
+                if (macmon_path := shutil.which("macmon")) is not None:
+                    tg.start_soon(self._monitor_macmon, macmon_path)
+                tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
+            tg.start_soon(self._watch_system_info)
+            tg.start_soon(self._monitor_memory_usage)
+            tg.start_soon(self._monitor_misc)
+
+            nc = await NodeConfig.gather()
+            if nc is not None:
+                await self.info_sender.send(nc)
+            sni = await StaticNodeInformation.gather()
+            await self.info_sender.send(sni)
+
+    def shutdown(self):
+        self._tg.cancel_scope.cancel()
+
+    async def _monitor_misc(self):
+        if self.misc_poll_interval is None:
+            return
+        prev = await MiscData.gather()
+        await self.info_sender.send(prev)
+        while True:
+            curr = await MiscData.gather()
+            if prev != curr:
+                prev = curr
+                await self.info_sender.send(curr)
+            await anyio.sleep(self.misc_poll_interval)
+
+    async def _monitor_system_profiler_thunderbolt_data(self):
+        if self.system_profiler_interval is None:
+            return
+        iface_map = await _gather_iface_map()
+        if iface_map is None:
+            return
+
+        old_idents = []
+        while True:
+            data = await ThunderboltConnectivity.gather()
+            assert data is not None
+
+            idents = [it for i in data if (it := i.ident(iface_map)) is not None]
+            if idents != old_idents:
+                await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
+            old_idents = idents
+
+            conns = [it for i in data if (it := i.conn()) is not None]
+            await self.info_sender.send(MacThunderboltConnections(conns=conns))
+
+            await anyio.sleep(self.system_profiler_interval)
+
+    async def _monitor_memory_usage(self):
+        override_memory_env = os.getenv("OVERRIDE_MEMORY_MB")
+        override_memory: int | None = (
+            Memory.from_mb(int(override_memory_env)).in_bytes
+            if override_memory_env
+            else None
+        )
+        if self.memory_poll_rate is None:
+            return
+        while True:
+            await self.info_sender.send(
+                MemoryUsage.from_psutil(override_memory=override_memory)
+            )
+            await anyio.sleep(self.memory_poll_rate)
+
+    async def _watch_system_info(self):
+        if self.interface_watcher_interval is None:
+            return
+        old_nics = []
+        while True:
+            nics = get_network_interfaces()
+            if nics != old_nics:
+                old_nics = nics
+                await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
+            await anyio.sleep(self.interface_watcher_interval)
+
+    async def _monitor_macmon(self, macmon_path: str):
+        if self.macmon_interval is None:
+            return
+        # macmon pipe --interval [interval in ms]
+        try:
+            async with await open_process(
+                [macmon_path, "pipe", "--interval", str(self.macmon_interval * 1000)]
+            ) as p:
+                if not p.stdout:
+                    logger.critical("MacMon closed stdout")
+                    return
+                async for text in TextReceiveStream(
+                    BufferedByteReceiveStream(p.stdout)
+                ):
+                    await self.info_sender.send(MacmonMetrics.from_raw_json(text))
+        except CalledProcessError as e:
+            stderr_msg = "no stderr"
+            stderr_output = cast(bytes | str | None, e.stderr)
+            if stderr_output is not None:
+                stderr_msg = (
+                    stderr_output.decode()
+                    if isinstance(stderr_output, bytes)
+                    else str(stderr_output)
+                )
+            logger.warning(
+                f"MacMon failed with return code {e.returncode}: {stderr_msg}"
+            )
--- a/src/exo/utils/info_gatherer/macmon.py
+++ b/src/exo/utils/info_gatherer/macmon.py
@@ -0,0 +1,70 @@
+from typing import Self
+
+from pydantic import BaseModel
+
+from exo.shared.types.profiling import MemoryUsage, SystemPerformanceProfile
+from exo.utils.pydantic_ext import TaggedModel
+
+
+class _TempMetrics(BaseModel, extra="ignore"):
+    """Temperature-related metrics returned by macmon."""
+
+    cpu_temp_avg: float
+    gpu_temp_avg: float
+
+
+class _MemoryMetrics(BaseModel, extra="ignore"):
+    """Memory-related metrics returned by macmon."""
+
+    ram_total: int
+    ram_usage: int
+    swap_total: int
+    swap_usage: int
+
+
+class RawMacmonMetrics(BaseModel, extra="ignore"):
+    """Complete set of metrics returned by macmon.
+
+    Unknown fields are ignored for forward-compatibility.
+    """
+
+    timestamp: str  # ignored
+    temp: _TempMetrics
+    memory: _MemoryMetrics
+    ecpu_usage: tuple[int, float]  # freq mhz, usage %
+    pcpu_usage: tuple[int, float]  # freq mhz, usage %
+    gpu_usage: tuple[int, float]  # freq mhz, usage %
+    all_power: float
+    ane_power: float
+    cpu_power: float
+    gpu_power: float
+    gpu_ram_power: float
+    ram_power: float
+    sys_power: float
+
+
+class MacmonMetrics(TaggedModel):
+    system_profile: SystemPerformanceProfile
+    memory: MemoryUsage
+
+    @classmethod
+    def from_raw(cls, raw: RawMacmonMetrics) -> Self:
+        return cls(
+            system_profile=SystemPerformanceProfile(
+                gpu_usage=raw.gpu_usage[1],
+                temp=raw.temp.gpu_temp_avg,
+                sys_power=raw.sys_power,
+                pcpu_usage=raw.pcpu_usage[1],
+                ecpu_usage=raw.ecpu_usage[1],
+            ),
+            memory=MemoryUsage.from_bytes(
+                ram_total=raw.memory.ram_total,
+                ram_available=(raw.memory.ram_total - raw.memory.ram_usage),
+                swap_total=raw.memory.swap_total,
+                swap_available=(raw.memory.swap_total - raw.memory.swap_usage),
+            ),
+        )
+
+    @classmethod
+    def from_raw_json(cls, json: str) -> Self:
+        return cls.from_raw(RawMacmonMetrics.model_validate_json(json))
--- a/src/exo/utils/info_gatherer/net_profile.py
+++ b/src/exo/utils/info_gatherer/net_profile.py
@@ -1,3 +1,5 @@
+from collections.abc import Mapping
+
 import anyio
 import httpx
 from anyio import create_task_group
@@ -5,6 +7,7 @@ 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

@@ -18,8 +21,9 @@ 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
-        target_ip = f"[{target_ip}]"
-    url = f"http://{target_ip}:52415/node_id"
+        url = f"http://[{target_ip}]:52415/node_id"
+    else:
+        url = f"http://{target_ip}:52415/node_id"

    remote_node_id = None
    last_error = None
@@ -73,7 +77,9 @@ async def check_reachability(


 async def check_reachable(
-    topology: Topology, self_node_id: NodeId
+    topology: Topology,
+    self_node_id: NodeId,
+    node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> dict[NodeId, set[str]]:
    """Check which nodes are reachable and return their IPs."""

@@ -91,16 +97,16 @@ async def check_reachable(
        httpx.AsyncClient(timeout=timeout, limits=limits) as client,
        create_task_group() as tg,
    ):
-        for node in topology.list_nodes():
-            if not node.node_profile:
+        for node_id in topology.list_nodes():
+            if node_id not in node_network:
                continue
-            if node.node_id == self_node_id:
+            if node_id == self_node_id:
                continue
-            for iface in node.node_profile.network_interfaces:
+            for iface in node_network[node_id].interfaces:
                tg.start_soon(
                    check_reachability,
                    iface.ip_address,
-                    node.node_id,
+                    node_id,
                    reachable,
                    client,
                )
--- a/src/exo/utils/info_gatherer/system_info.py
+++ b/src/exo/utils/info_gatherer/system_info.py
--- a/src/exo/utils/info_gatherer/tests/test_tb_parsing.py
+++ b/src/exo/utils/info_gatherer/tests/test_tb_parsing.py
@@ -0,0 +1,24 @@
+import sys
+
+import pytest
+
+from exo.shared.types.thunderbolt import (
+    ThunderboltConnectivity,
+)
+from exo.utils.info_gatherer.info_gatherer import (
+    _gather_iface_map,  # pyright: ignore[reportPrivateUsage]
+)
+
+
+@pytest.mark.anyio
+@pytest.mark.skipif(
+    sys.platform != "darwin", reason="Thunderbolt info can only be gathered on macos"
+)
+async def test_tb_parsing():
+    data = await ThunderboltConnectivity.gather()
+    ifaces = await _gather_iface_map()
+    assert ifaces
+    assert data
+    for datum in data:
+        datum.ident(ifaces)
+        datum.conn()
--- a/src/exo/utils/pydantic_ext.py
+++ b/src/exo/utils/pydantic_ext.py
@@ -19,11 +19,20 @@ class CamelCaseModel(BaseModel):
        alias_generator=to_camel,
        validate_by_name=True,
        extra="forbid",
-        # I want to reenable this ASAP, but it's causing an issue with TaskStatus
        strict=True,
    )


+class FrozenModel(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        validate_by_name=True,
+        extra="forbid",
+        strict=True,
+        frozen=True,
+    )
+
+
 class TaggedModel(CamelCaseModel):
    @model_serializer(mode="wrap")
    def _serialize(self, handler: SerializerFunctionWrapHandler):
--- a/src/exo/utils/tests/test_mp_channel.py
+++ b/src/exo/utils/tests/test_mp_channel.py
@@ -28,9 +28,8 @@ def bar(send: MpSender[str]):
    send.close()


-# not async, just want the fail_after
@pytest.mark.anyio
-async def test_channel_setup():
+async def test_channel_ipc():
    with fail_after(0.5):
        s, r = mp_channel[str]()
        p1 = mp.Process(target=foo, args=(r,))
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -5,6 +5,7 @@ import shutil
 import ssl
 import time
 import traceback
+from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -525,7 +526,7 @@ async def download_progress_for_local_path(

 async def download_shard(
    shard: ShardMetadata,
-    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
+    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    max_parallel_downloads: int = 8,
    skip_download: bool = False,
    allow_patterns: list[str] | None = None,
@@ -566,9 +567,9 @@ async def download_shard(
    )
    file_progress: dict[str, RepoFileDownloadProgress] = {}

-    def on_progress_wrapper(
+    async def on_progress_wrapper(
        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ):
+    ) -> None:
        start_time = (
            file_progress[file.path].start_time
            if file.path in file_progress
@@ -604,7 +605,7 @@ async def download_shard(
            else "in_progress",
            start_time=start_time,
        )
-        on_progress(
+        await on_progress(
            shard,
            calculate_repo_progress(
                shard,
@@ -632,14 +633,21 @@ async def download_shard(

    semaphore = asyncio.Semaphore(max_parallel_downloads)

-    async def download_with_semaphore(file: FileListEntry):
+    def schedule_progress(
+        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
+    ) -> None:
+        asyncio.create_task(
+            on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed)
+        )
+
+    async def download_with_semaphore(file: FileListEntry) -> None:
        async with semaphore:
            await download_file_with_retry(
                str(shard.model_meta.model_id),
                revision,
                file.path,
                target_dir,
-                lambda curr_bytes, total_bytes, is_renamed: on_progress_wrapper(
+                lambda curr_bytes, total_bytes, is_renamed: schedule_progress(
                    file, curr_bytes, total_bytes, is_renamed
                ),
            )
@@ -651,7 +659,7 @@ async def download_shard(
    final_repo_progress = calculate_repo_progress(
        shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
    )
-    on_progress(shard, final_repo_progress)
+    await on_progress(shard, final_repo_progress)
    if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
        return target_dir / gguf.path, final_repo_progress
    else:
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -1,4 +1,5 @@
 import asyncio
+from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable

@@ -48,7 +49,8 @@ class SingletonShardDownloader(ShardDownloader):
        self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        self.shard_downloader.on_progress(callback)

@@ -83,7 +85,8 @@ class CachedShardDownloader(ShardDownloader):
        self.cache: dict[tuple[str, ShardMetadata], Path] = {}

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        self.shard_downloader.on_progress(callback)

@@ -113,17 +116,18 @@ class ResumableShardDownloader(ShardDownloader):
    def __init__(self, max_parallel_downloads: int = 8):
        self.max_parallel_downloads = max_parallel_downloads
        self.on_progress_callbacks: list[
-            Callable[[ShardMetadata, RepoDownloadProgress], None]
+            Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
        ] = []

-    def on_progress_wrapper(
+    async def on_progress_wrapper(
        self, shard: ShardMetadata, progress: RepoDownloadProgress
    ) -> None:
        for callback in self.on_progress_callbacks:
-            callback(shard, progress)
+            await callback(shard, progress)

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        self.on_progress_callbacks.append(callback)

--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections.abc import Awaitable
 from copy import copy
 from datetime import timedelta
 from pathlib import Path
@@ -31,7 +32,8 @@ class ShardDownloader(ABC):

    @abstractmethod
    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        pass

@@ -59,7 +61,8 @@ class NoopShardDownloader(ShardDownloader):
        return Path("/tmp/noop_shard")

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        pass

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

    def __init__(self, original_layer: _LayerCallable):
        super().__init__()
-        # Set twice to avoid __setattr__ recursion
        object.__setattr__(self, "_original_layer", original_layer)
-        self.original_layer: _LayerCallable = original_layer
+
+    @property
+    def original_layer(self) -> _LayerCallable:
+        return cast(_LayerCallable, object.__getattribute__(self, "_original_layer"))

    # Calls __getattr__ for any attributes not found on nn.Module (e.g. use_sliding)
    if not TYPE_CHECKING:
@@ -58,7 +60,7 @@ class CustomMlxLayer(nn.Module):
                return super().__getattr__(name)
            except AttributeError:
                original_layer = object.__getattribute__(self, "_original_layer")
-                return object.__getattribute__(original_layer, name)
+                return getattr(original_layer, name)


 class PipelineFirstLayer(CustomMlxLayer):
@@ -106,7 +108,6 @@ class PipelineLastLayer(CustomMlxLayer):
            if cache is not None:
                cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]

-        output = mx.distributed.all_gather(output, group=self.group)[-output.shape[0] :]
        return output


@@ -168,11 +169,21 @@ def pipeline_auto_parallel(
        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
            start_layer:end_layer
        ]
-        inner_model_instance.swa_idx = inner_model_instance.layer_types.index(  # type: ignore
-            "sliding_attention"
+        # We can assume the model has at least one layer thanks to placement.
+        # If a layer type doesn't exist, we can set it to 0.
+        inner_model_instance.swa_idx = (
+            0
+            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
+            else inner_model_instance.layer_types.index(  # type: ignore
+                "sliding_attention"
+            )
        )
-        inner_model_instance.ga_idx = inner_model_instance.layer_types.index(  # type: ignore
-            "full_attention"
+        inner_model_instance.ga_idx = (
+            0
+            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
+            else inner_model_instance.layer_types.index(  # type: ignore
+                "full_attention"
+            )
        )

    _set_layers(model, layers)
@@ -181,6 +192,36 @@ 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


--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -8,12 +8,13 @@ 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.openai_responses import ResponsesRequest
+from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.runner_response import (
    GenerationResponse,
 )
@@ -52,9 +53,14 @@ def warmup_inference(

    warmup_prompt = apply_chat_template(
        tokenizer=tokenizer,
-        task_params=ResponsesRequest(
+        chat_task_data=ChatCompletionTaskParams(
            model="",
-            input=content,
+            messages=[
+                ChatCompletionMessage(
+                    role="user",
+                    content=content,
+                )
+            ],
        ),
    )

@@ -109,68 +115,15 @@ 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: ResponsesRequest,
-    is_bench: bool = False,
+    task: ChatCompletionTaskParams,
+    prompt: str,
 ) -> 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}")
@@ -178,11 +131,6 @@ 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]] = []
@@ -194,20 +142,9 @@ 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,
    )

-    # 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 = ""
-
+    max_tokens = task.max_tokens or MAX_TOKENS
    for out in stream_generate(
        model=model,
        tokenizer=tokenizer,
@@ -216,35 +153,15 @@ 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 is_done:
+        if out.finish_reason is not None:
            stats = GenerationStats(
                prompt_tps=float(out.prompt_tps),
                generation_tps=float(out.generation_tps),
@@ -252,33 +169,22 @@ def mlx_generate(
                generation_tokens=int(out.generation_tokens),
                peak_memory_usage=Memory.from_gb(out.peak_memory),
            )
-            if not stop_matched and out.finish_reason not in get_args(FinishReason):
+
+            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
                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=text,
+            text=out.text,
            token=out.token,
-            logprob=logprob,
-            top_logprobs=top_logprobs,
-            finish_reason=finish_reason,
+            finish_reason=cast(FinishReason | None, out.finish_reason),
            stats=stats,
        )

-        if is_done:
+        if out.finish_reason is not None:
            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?
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -42,9 +42,10 @@ 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.openai_responses import ResponsesRequest
+from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.instances import (
    BoundInstance,
    MlxJacclInstance,
@@ -185,26 +186,37 @@ def mlx_distributed_init(
                group = mx.distributed.init(backend="ring", strict=True)

            case MlxJacclInstance(
-                ibv_devices=ibv_devices, jaccl_coordinators=jaccl_coordinators
+                jaccl_devices=jaccl_devices, jaccl_coordinators=jaccl_coordinators
            ):
+                assert all(
+                    jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
+                )
                # Use RDMA connectivity matrix
                coordination_file = (
                    f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
                )
-                ibv_devices_json = json.dumps(ibv_devices)
+                jaccl_devices_json = json.dumps(jaccl_devices)

                with open(coordination_file, "w") as f:
-                    _ = f.write(ibv_devices_json)
+                    _ = f.write(jaccl_devices_json)

                jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]

-                logger.info(f"rank {rank} MLX_IBV_DEVICES: {ibv_devices_json}")
+                # TODO: update once upstream fixes
+                logger.info(
+                    f"rank {rank} MLX_IBV_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
+                )
                logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
                os.environ["MLX_IBV_DEVICES"] = coordination_file
                os.environ["MLX_RANK"] = str(rank)
                os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                group = mx.distributed.init(backend="jaccl", strict=True)

+            case _:
+                raise ValueError(
+                    f"Unsupported instance type for MLX distributed: {type(bound_instance.instance)}"
+                )
+
        logger.info(f"Rank {rank} mlx distributed initialization complete")

        return group
@@ -389,59 +401,51 @@ def load_tokenizer_for_model_id(model_id: str, model_path: Path) -> TokenizerWra

 def apply_chat_template(
    tokenizer: TokenizerWrapper,
-    task_params: ResponsesRequest,
+    chat_task_data: ChatCompletionTaskParams,
 ) -> str:
-    """Convert ResponsesRequest to a chat template prompt.
+    # Now we can properly access the messages
+    messages = chat_task_data.messages

-    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]] = []
-
-    # Add system message (instructions) if present
-    if task_params.instructions:
-        formatted_messages.append(
-            {"role": "system", "content": task_params.instructions}
-        )
-
-    # 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")
+    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
-            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,
+            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
+        formatted_messages.append(
+            {k: v for k, v in message.model_dump().items() if v is not None}  # type: ignore
        )

+    prompt: str = tokenizer.apply_chat_template(
+        formatted_messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        tools=chat_task_data.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.
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -16,8 +16,7 @@ from exo.shared.types.events import (
    ForwarderEvent,
    IndexedEvent,
    NodeDownloadProgress,
-    NodeMemoryMeasured,
-    NodePerformanceMeasured,
+    NodeGatheredInfo,
    TaskCreated,
    TaskStatusUpdated,
    TopologyEdgeCreated,
@@ -25,7 +24,6 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.models import ModelId
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import MemoryPerformanceProfile, NodePerformanceProfile
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
    CreateRunner,
@@ -34,7 +32,7 @@ from exo.shared.types.tasks import (
    Task,
    TaskStatus,
 )
-from exo.shared.types.topology import Connection
+from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import (
    DownloadCompleted,
    DownloadOngoing,
@@ -45,14 +43,14 @@ from exo.shared.types.worker.runners import RunnerId
 from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import OrderedBuffer
+from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
+from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.worker.download.download_utils import (
    map_repo_download_progress_to_download_progress_data,
 )
 from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDownloader
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
-from exo.worker.utils import start_polling_memory_metrics, start_polling_node_metrics
-from exo.worker.utils.net_profile import check_reachable


 class Worker:
@@ -86,7 +84,7 @@ class Worker:
        self.state: State = State()
        self.download_status: dict[ModelId, DownloadProgress] = {}
        self.runners: dict[RunnerId, RunnerSupervisor] = {}
-        self._tg: TaskGroup | None = None
+        self._tg: TaskGroup = create_task_group()

        self._nack_cancel_scope: CancelScope | None = None
        self._nack_attempts: int = 0
@@ -98,37 +96,13 @@ class Worker:
    async def run(self):
        logger.info("Starting Worker")

-        # TODO: CLEANUP HEADER
-        async def resource_monitor_callback(
-            node_performance_profile: NodePerformanceProfile,
-        ) -> None:
-            await self.event_sender.send(
-                NodePerformanceMeasured(
-                    node_id=self.node_id,
-                    node_profile=node_performance_profile,
-                    when=str(datetime.now(tz=timezone.utc)),
-                ),
-            )
+        info_send, info_recv = channel[GatheredInfo]()
+        info_gatherer: InfoGatherer = InfoGatherer(info_send)

-        async def memory_monitor_callback(
-            memory_profile: MemoryPerformanceProfile,
-        ) -> None:
-            await self.event_sender.send(
-                NodeMemoryMeasured(
-                    node_id=self.node_id,
-                    memory=memory_profile,
-                    when=str(datetime.now(tz=timezone.utc)),
-                )
-            )
-
-        # END CLEANUP
-
-        async with create_task_group() as tg:
-            self._tg = tg
+        async with self._tg as tg:
+            tg.start_soon(info_gatherer.run)
+            tg.start_soon(self._forward_info, info_recv)
            tg.start_soon(self.plan_step)
-            tg.start_soon(start_polling_node_metrics, resource_monitor_callback)
-
-            tg.start_soon(start_polling_memory_metrics, memory_monitor_callback)
            tg.start_soon(self._emit_existing_download_progress)
            tg.start_soon(self._connection_message_event_writer)
            tg.start_soon(self._resend_out_for_delivery)
@@ -142,6 +116,17 @@ class Worker:
        for runner in self.runners.values():
            runner.shutdown()

+    async def _forward_info(self, recv: Receiver[GatheredInfo]):
+        with recv as info_stream:
+            async for info in info_stream:
+                await self.event_sender.send(
+                    NodeGatheredInfo(
+                        node_id=self.node_id,
+                        when=str(datetime.now(tz=timezone.utc)),
+                        info=info,
+                    )
+                )
+
    async def _event_applier(self):
        with self.global_event_receiver as events:
            async for f_event in events:
@@ -161,7 +146,6 @@ class Worker:
                    self._nack_cancel_scope is None
                    or self._nack_cancel_scope.cancel_called
                ):
-                    assert self._tg
                    # Request the next index.
                    self._tg.start_soon(
                        self._nack_request, self.state.last_event_applied_idx + 1
@@ -252,8 +236,7 @@ class Worker:
                    await self.runners[self._task_to_runner_id(task)].start_task(task)

    def shutdown(self):
-        if self._tg:
-            self._tg.cancel_scope.cancel()
+        self._tg.cancel_scope.cancel()

    def _task_to_runner_id(self, task: Task):
        instance = self.state.instances[task.instance_id]
@@ -270,24 +253,28 @@ class Worker:
        match msg.connection_type:
            case ConnectionMessageType.Connected:
                return TopologyEdgeCreated(
-                    edge=Connection(
-                        local_node_id=self.node_id,
-                        send_back_node_id=msg.node_id,
-                        send_back_multiaddr=Multiaddr(
-                            address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
                        ),
-                    )
+                    ),
                )

            case ConnectionMessageType.Disconnected:
                return TopologyEdgeDeleted(
-                    edge=Connection(
-                        local_node_id=self.node_id,
-                        send_back_node_id=msg.node_id,
-                        send_back_multiaddr=Multiaddr(
-                            address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
                        ),
-                    )
+                    ),
                )

    async def _nack_request(self, since_idx: int) -> None:
@@ -336,7 +323,6 @@ class Worker:
            event_sender=self.event_sender.clone(),
        )
        self.runners[task.bound_instance.bound_runner_id] = runner
-        assert self._tg
        self._tg.start_soon(runner.run)
        return runner

@@ -359,8 +345,7 @@ class Worker:
        last_progress_time = 0.0
        throttle_interval_secs = 1.0

-        # TODO: i hate callbacks
-        def download_progress_callback(
+        async def download_progress_callback(
            shard: ShardMetadata, progress: RepoDownloadProgress
        ) -> None:
            nonlocal self
@@ -372,11 +357,10 @@ class Worker:
                    total_bytes=progress.total_bytes,
                )
                self.download_status[shard.model_meta.model_id] = status
-                # Footgun!
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                    NodeDownloadProgress(download_progress=status)
                )
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                    TaskStatusUpdated(
                        task_id=task.task_id, task_status=TaskStatus.Complete
                    )
@@ -393,13 +377,12 @@ class Worker:
                    ),
                )
                self.download_status[shard.model_meta.model_id] = status
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                    NodeDownloadProgress(download_progress=status)
                )
                last_progress_time = current_time()

        self.shard_downloader.on_progress(download_progress_callback)
-        assert self._tg
        self._tg.start_soon(self.shard_downloader.ensure_shard, task.shard_metadata)

    async def _forward_events(self) -> None:
@@ -420,9 +403,14 @@ class Worker:

    async def _poll_connection_updates(self):
        while True:
-            # TODO: EdgeDeleted
-            edges = set(self.state.topology.list_connections())
-            conns = await check_reachable(self.state.topology, self.node_id)
+            edges = set(
+                conn.edge for conn in self.state.topology.out_edges(self.node_id)
+            )
+            conns = await check_reachable(
+                self.state.topology,
+                self.node_id,
+                self.state.node_network,
+            )
            for nid in conns:
                for ip in conns[nid]:
                    if "127.0.0.1" in ip or "localhost" in ip:
@@ -430,26 +418,33 @@ class Worker:
                            f"Loopback connection should not happen: {ip=} for {nid=}"
                        )

-                    edge = Connection(
-                        local_node_id=self.node_id,
-                        send_back_node_id=nid,
+                    edge = SocketConnection(
                        # nonsense multiaddr
-                        send_back_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
+                        sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
                        if "." in ip
                        # nonsense multiaddr
                        else Multiaddr(address=f"/ip6/{ip}/tcp/52415"),
                    )
                    if edge not in edges:
                        logger.debug(f"ping discovered {edge=}")
-                        await self.event_sender.send(TopologyEdgeCreated(edge=edge))
+                        await self.event_sender.send(
+                            TopologyEdgeCreated(
+                                conn=Connection(
+                                    source=self.node_id, sink=nid, edge=edge
+                                )
+                            )
+                        )

-            for nid, conn in self.state.topology.out_edges(self.node_id):
+            for conn in self.state.topology.out_edges(self.node_id):
+                if not isinstance(conn.edge, SocketConnection):
+                    continue
                if (
-                    nid not in conns
-                    or conn.send_back_multiaddr.ip_address not in conns.get(nid, set())
+                    conn.sink not in conns
+                    or conn.edge.sink_multiaddr.ip_address
+                    not in conns.get(conn.sink, set())
                ):
                    logger.debug(f"ping failed to discover {conn=}")
-                    await self.event_sender.send(TopologyEdgeDeleted(edge=conn))
+                    await self.event_sender.send(TopologyEdgeDeleted(conn=conn))

            await anyio.sleep(10)

--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -21,7 +21,11 @@ from exo.shared.types.worker.downloads import (
    DownloadOngoing,
    DownloadProgress,
 )
-from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId
+from exo.shared.types.worker.instances import (
+    BoundInstance,
+    Instance,
+    InstanceId,
+)
 from exo.shared.types.worker.runners import (
    RunnerConnected,
    RunnerConnecting,
@@ -50,6 +54,16 @@ def plan(
    all_runners: Mapping[RunnerId, RunnerStatus],  # all global
    tasks: Mapping[TaskId, Task],
 ) -> Task | None:
+    from exo.plugins.registry import PluginRegistry
+
+    registry = PluginRegistry.get()
+
+    # Check plugin tasks first
+    for plugin in registry.all_plugins():
+        task = plugin.plan_task(runners, instances)
+        if task is not None:
+            return task
+
    # Python short circuiting OR logic should evaluate these sequentially.
    return (
        _kill_runner(runners, all_runners, instances)
@@ -113,7 +127,18 @@ def _model_needs_download(
    runners: Mapping[RunnerId, RunnerSupervisor],
    download_status: Mapping[ModelId, DownloadProgress],
 ) -> DownloadModel | None:
+    from exo.plugins.registry import PluginRegistry
+
+    registry = PluginRegistry.get()
+
    for runner in runners.values():
+        instance = runner.bound_instance.instance
+
+        # Check if any plugin wants to skip download for this instance
+        plugin = registry.get_plugin_for_instance(instance)
+        if plugin is not None and plugin.should_skip_download(instance):
+            continue
+
        model_id = runner.bound_instance.bound_shard.model_meta.model_id
        if isinstance(runner.status, RunnerIdle) and (
            model_id not in download_status
--- a/src/exo/worker/runner/bootstrap.py
+++ b/src/exo/worker/runner/bootstrap.py
@@ -4,7 +4,10 @@ import loguru

 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task
-from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
+from exo.shared.types.worker.instances import (
+    BoundInstance,
+    MlxJacclInstance,
+)
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender

@@ -17,12 +20,13 @@ def entrypoint(
    task_receiver: MpReceiver[Task],
    _logger: "loguru.Logger",
 ) -> None:
+    # Set FAST_SYNCH based on env var or JACCL device count
    fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
    if fast_synch_override == "on" or (
        fast_synch_override != "off"
        and (
            isinstance(bound_instance.instance, MlxJacclInstance)
-            and len(bound_instance.instance.ibv_devices) >= 2
+            and len(bound_instance.instance.jaccl_devices) >= 2
        )
    ):
        os.environ["MLX_METAL_FAST_SYNCH"] = "1"
@@ -34,11 +38,26 @@ def entrypoint(

    logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")

-    # Import main after setting global logger - this lets us just import logger from this module
+    # Route based on instance type (plugins or default MLX)
    try:
-        from exo.worker.runner.runner import main
+        from exo.plugins.registry import PluginRegistry, discover_plugins

-        main(bound_instance, event_sender, task_receiver)
+        # Discover plugins in subprocess (they aren't inherited from main process)
+        discover_plugins()
+
+        registry = PluginRegistry.get()
+        instance = bound_instance.instance
+
+        # Check if a plugin handles this instance type
+        plugin = registry.get_plugin_for_instance(instance)
+        if plugin is not None:
+            # Delegate to plugin runner
+            plugin.create_runner(bound_instance, event_sender, task_receiver)
+        else:
+            # MLX runner (default)
+            from exo.worker.runner.runner import main
+
+            main(bound_instance, event_sender, task_receiver)
    except ClosedResourceError:
        logger.warning("Runner communication closed unexpectedly")
    except Exception as e:
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -4,6 +4,7 @@ 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,
@@ -11,6 +12,7 @@ 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,
@@ -19,7 +21,6 @@ from exo.shared.types.events import (
    TaskAcknowledged,
    TaskStatusUpdated,
 )
-from exo.shared.types.openai_responses import ResponsesRequest
 from exo.shared.types.tasks import (
    ChatCompletion,
    ConnectToGroup,
@@ -50,6 +51,8 @@ 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,
@@ -172,21 +175,33 @@ def main(
                    )
                    assert model
                    assert tokenizer
+                    assert task_params.messages[0].content is not None

                    try:
-                        _check_for_debug_prompts(task_params)
+                        _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)

                        # 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:
@@ -201,8 +216,6 @@ def main(
                                                    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,
                                                ),
@@ -294,28 +307,44 @@ 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(task_params: ResponsesRequest) -> None:
-    """Check for debug prompt triggers in the input.
-
-    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")
+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 = task_params.input[0].content
+        prompt = prompt[0]

-    if not prompt:
-        return
+    if isinstance(prompt, ChatCompletionMessageText):
+        prompt = prompt.text

    if EXO_RUNNER_MUST_FAIL in prompt:
        logger.info("raising exception")
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -0,0 +1,201 @@
+# 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.types.api import ChatCompletionMessage
+from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
+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_meta=ModelMetadata(
+                model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
+                pretty_name="GPT-OSS 20B",
+                storage_size=Memory.from_gb(12),
+                n_layers=24,
+                hidden_size=2880,
+                supports_tensor=False,
+            ),
+            device_rank=rank,
+            world_size=world_size,
+            start_layer=start_layer,
+            end_layer=end_layer,
+            n_layers=24,
+        )
+
+        model, 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_meta=ModelMetadata(
+                model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
+                pretty_name="GPT-OSS 20B",
+                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]
--- a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
@@ -0,0 +1,146 @@
+import json
+import multiprocessing as mp
+import os
+import tempfile
+from typing import Any
+
+import mlx.core as mx
+import mlx.nn as mlx_nn
+import pytest
+
+from exo.worker.engines.mlx.auto_parallel import (
+    CustomMlxLayer,
+    PipelineFirstLayer,
+    PipelineLastLayer,
+    patch_pipeline_model,
+)
+from exo.worker.tests.unittests.test_mlx.conftest import MockLayer
+
+
+def run_pipeline_device(
+    rank: int,
+    world_size: int,
+    hostfile_path: str,
+    result_queue: Any,  # pyright: ignore[reportAny]
+) -> None:
+    import os
+
+    os.environ["MLX_HOSTFILE"] = hostfile_path
+    os.environ["MLX_RANK"] = str(rank)
+
+    class MockLayerInner(mlx_nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.custom_attr = "test_value"
+
+        def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+            return x * 2
+
+    class MockModel(mlx_nn.Module):
+        def __init__(self, layers: list[mlx_nn.Module]) -> None:
+            super().__init__()
+            self.layers = layers
+
+        def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+            for layer in self.layers:
+                x = layer(x, *args, **kwargs)  # pyright: ignore[reportUnknownVariableType]
+            return x  # pyright: ignore[reportUnknownVariableType]
+
+    try:
+        group = mx.distributed.init(backend="ring", strict=True)
+
+        mock = MockLayerInner()
+        first = PipelineFirstLayer(mock, r=rank, group=group)
+        composed = PipelineLastLayer(first, r=rank, s=world_size, group=group)
+
+        # Wrap in a mock model, then wrap in PipelineParallelModel for all_gather
+        inner_model = MockModel([composed])
+        model = patch_pipeline_model(inner_model, group)
+
+        x = mx.ones((1, 4))
+        result = model(x)
+        mx.eval(result)
+        success = result.shape == x.shape
+        result_queue.put((rank, success, result))  # pyright: ignore[reportAny]
+    except Exception as e:
+        result_queue.put((rank, False, str(e)))  # pyright: ignore[reportAny]
+
+
+def test_single_wrapper_delegates_attributes() -> None:
+    mock = MockLayer()
+    wrapped = CustomMlxLayer(mock)
+
+    assert wrapped.custom_attr == "test_value"  # type: ignore[attr-defined]
+    assert wrapped.use_sliding is True  # type: ignore[attr-defined]
+
+
+def test_composed_wrappers_delegate_attributes() -> None:
+    mock = MockLayer()
+    group = mx.distributed.init()
+
+    first = PipelineFirstLayer(mock, r=0, group=group)
+    composed = PipelineLastLayer(first, r=0, s=1, group=group)
+
+    assert composed.custom_attr == "test_value"  # type: ignore[attr-defined]
+    assert composed.use_sliding is True  # type: ignore[attr-defined]
+
+
+def test_missing_attribute_raises() -> None:
+    mock = MockLayer()
+    wrapped = CustomMlxLayer(mock)
+
+    with pytest.raises(AttributeError):
+        _ = wrapped.nonexistent_attr  # type: ignore[attr-defined]
+
+
+def test_composed_call_works() -> None:
+    ctx = mp.get_context("spawn")
+
+    world_size = 2
+    base_port = 29500
+
+    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        hostfile_path = f.name
+
+    try:
+        result_queue: Any = ctx.Queue()
+
+        processes: list[Any] = []
+        for rank in range(world_size):
+            p = ctx.Process(
+                target=run_pipeline_device,
+                args=(rank, world_size, hostfile_path, result_queue),
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:  # pyright: ignore[reportAny]
+            p.join(timeout=10)  # pyright: ignore[reportAny]
+
+        results: dict[int, Any] = {}
+        errors: dict[int, str] = {}
+        while not result_queue.empty():  # pyright: ignore[reportAny]
+            rank, success, value = result_queue.get()  # pyright: ignore[reportAny]
+            if success:
+                results[rank] = value
+            else:
+                errors[rank] = value
+
+        assert len(results) == world_size, (
+            f"Expected {world_size} results, got {len(results)}. Errors: {errors}"
+        )
+
+        for rank in range(world_size):
+            assert rank in results, (
+                f"Device {rank} failed: {errors.get(rank, 'unknown')}"
+            )
+            result_array = results[rank]
+            # Both devices see the final result (4.0) after all_gather
+            assert (result_array == 4.0).all(), (
+                f"Device {rank}: expected 4.0, got {result_array}"
+            )
+    finally:
+        os.unlink(hostfile_path)
--- a/src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py
@@ -0,0 +1,230 @@
+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}"
--- a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
+++ b/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.openai_responses import ResponsesRequest
+from exo.shared.types.api import ChatCompletionTaskParams
 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=ResponsesRequest(model=MODEL_A_ID, input=""),
+        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
    )

    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=ResponsesRequest(model=MODEL_A_ID, input=""),
+        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
    )

    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=ResponsesRequest(model=MODEL_A_ID, input=""),
+        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
    )

    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=ResponsesRequest(model=MODEL_A_ID, input=""),
+        task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
    )

    other_task_id = TaskId("other-task")
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -5,6 +5,7 @@ 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,
@@ -13,9 +14,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,
@@ -84,11 +85,11 @@ SHUTDOWN_TASK = Shutdown(
    runner_id=RUNNER_1_ID,
 )

-CHAT_PARAMS = ResponsesRequest(
+CHAT_PARAMS = ChatCompletionTaskParams(
    model=str(MODEL_A_ID),
-    input="hello",
+    messages=[ChatCompletionMessage(role="user", content="hello")],
    stream=True,
-    max_output_tokens=4,
+    max_tokens=4,
    temperature=0.0,
 )

@@ -113,6 +114,10 @@ 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")
--- a/src/exo/worker/utils/init.py
+++ b/src/exo/worker/utils/init.py
@@ -1,6 +0,0 @@
-from .profile import start_polling_memory_metrics, start_polling_node_metrics
-
-__all__ = [
-    "start_polling_node_metrics",
-    "start_polling_memory_metrics",
-]
--- a/src/exo/worker/utils/macmon.py
+++ b/src/exo/worker/utils/macmon.py
@@ -1,103 +0,0 @@
-import platform
-import shutil
-from subprocess import CalledProcessError
-from typing import cast
-
-from anyio import run_process
-from pydantic import BaseModel, ConfigDict, ValidationError
-
-
-class MacMonError(Exception):
-    """Exception raised for errors in the MacMon functions."""
-
-
-def _get_binary_path() -> str:
-    """
-    Get the path to the macmon binary.
-
-    Raises:
-        MacMonError: If the binary doesn't exist or can't be made executable.
-    """
-    # Check for macOS with ARM chip
-    system = platform.system().lower()
-    machine = platform.machine().lower()
-
-    if system != "darwin" or not (
-        "arm" in machine or "m1" in machine or "m2" in machine
-    ):
-        raise MacMonError("MacMon only supports macOS with Apple Silicon (ARM) chips")
-
-    path = shutil.which("macmon")
-
-    if path is None:
-        raise MacMonError("MacMon not found in PATH")
-
-    return path
-
-
-class TempMetrics(BaseModel):
-    """Temperature-related metrics returned by macmon."""
-
-    cpu_temp_avg: float
-    gpu_temp_avg: float
-
-    model_config = ConfigDict(extra="ignore")
-
-
-class Metrics(BaseModel):
-    """Complete set of metrics returned by macmon.
-
-    Unknown fields are ignored for forward-compatibility.
-    """
-
-    all_power: float
-    ane_power: float
-    cpu_power: float
-    ecpu_usage: tuple[int, float]
-    gpu_power: float
-    gpu_ram_power: float
-    gpu_usage: tuple[int, float]
-    pcpu_usage: tuple[int, float]
-    ram_power: float
-    sys_power: float
-    temp: TempMetrics
-    timestamp: str
-
-    model_config = ConfigDict(extra="ignore")
-
-
-async def get_metrics_async() -> Metrics:
-    """
-    Asynchronously run the binary and return the metrics as a Python dictionary.
-
-    Args:
-        binary_path: Optional path to the binary. If not provided, will use the bundled binary.
-
-    Returns:
-        A mapping containing system metrics.
-
-    Raises:
-        MacMonError: If there's an error running the binary.
-    """
-    path = _get_binary_path()
-
-    try:
-        # TODO: Keep Macmon running in the background?
-        result = await run_process([path, "pipe", "-s", "1"])
-
-        return Metrics.model_validate_json(result.stdout.decode().strip())
-
-    except ValidationError as e:
-        raise MacMonError(f"Error parsing JSON output: {e}") from e
-    except CalledProcessError as e:
-        stderr_msg = "no stderr"
-        stderr_output = cast(bytes | str | None, e.stderr)
-        if stderr_output is not None:
-            stderr_msg = (
-                stderr_output.decode()
-                if isinstance(stderr_output, bytes)
-                else str(stderr_output)
-            )
-        raise MacMonError(
-            f"MacMon failed with return code {e.returncode}: {stderr_msg}"
-        ) from e
--- a/src/exo/worker/utils/profile.py
+++ b/src/exo/worker/utils/profile.py
@@ -1,114 +0,0 @@
-import asyncio
-import os
-import platform
-from typing import Any, Callable, Coroutine
-
-import anyio
-from loguru import logger
-
-from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import (
-    MemoryPerformanceProfile,
-    NodePerformanceProfile,
-    SystemPerformanceProfile,
-)
-
-from .macmon import (
-    MacMonError,
-    Metrics,
-)
-from .macmon import (
-    get_metrics_async as macmon_get_metrics_async,
-)
-from .system_info import (
-    get_friendly_name,
-    get_model_and_chip,
-    get_network_interfaces,
-)
-
-
-async def get_metrics_async() -> Metrics | None:
-    """Return detailed Metrics on macOS or a minimal fallback elsewhere."""
-
-    if platform.system().lower() == "darwin":
-        return await macmon_get_metrics_async()
-
-
-def get_memory_profile() -> MemoryPerformanceProfile:
-    """Construct a MemoryPerformanceProfile using psutil"""
-    override_memory_env = os.getenv("OVERRIDE_MEMORY_MB")
-    override_memory: int | None = (
-        Memory.from_mb(int(override_memory_env)).in_bytes
-        if override_memory_env
-        else None
-    )
-
-    return MemoryPerformanceProfile.from_psutil(override_memory=override_memory)
-
-
-async def start_polling_memory_metrics(
-    callback: Callable[[MemoryPerformanceProfile], Coroutine[Any, Any, None]],
-    *,
-    poll_interval_s: float = 0.5,
-) -> None:
-    """Continuously poll and emit memory-only metrics at a faster cadence.
-
-    Parameters
-    - callback: coroutine called with a fresh MemoryPerformanceProfile each tick
-    - poll_interval_s: interval between polls
-    """
-    while True:
-        try:
-            mem = get_memory_profile()
-            await callback(mem)
-        except MacMonError as e:
-            logger.opt(exception=e).error("Memory Monitor encountered error")
-        finally:
-            await anyio.sleep(poll_interval_s)
-
-
-async def start_polling_node_metrics(
-    callback: Callable[[NodePerformanceProfile], Coroutine[Any, Any, None]],
-):
-    poll_interval_s = 1.0
-    while True:
-        try:
-            metrics = await get_metrics_async()
-            if metrics is None:
-                return
-
-            network_interfaces = get_network_interfaces()
-            # these awaits could be joined but realistically they should be cached
-            model_id, chip_id = await get_model_and_chip()
-            friendly_name = await get_friendly_name()
-
-            # do the memory profile last to get a fresh reading to not conflict with the other memory profiling loop
-            memory_profile = get_memory_profile()
-
-            await callback(
-                NodePerformanceProfile(
-                    model_id=model_id,
-                    chip_id=chip_id,
-                    friendly_name=friendly_name,
-                    network_interfaces=network_interfaces,
-                    memory=memory_profile,
-                    system=SystemPerformanceProfile(
-                        gpu_usage=metrics.gpu_usage[1],
-                        temp=metrics.temp.gpu_temp_avg,
-                        sys_power=metrics.sys_power,
-                        pcpu_usage=metrics.pcpu_usage[1],
-                        ecpu_usage=metrics.ecpu_usage[1],
-                        ane_power=metrics.ane_power,
-                    ),
-                )
-            )
-
-        except asyncio.TimeoutError:
-            logger.warning(
-                "[resource_monitor] Operation timed out after 30s, skipping this cycle."
-            )
-        except MacMonError as e:
-            logger.opt(exception=e).error("Resource Monitor encountered error")
-            return
-        finally:
-            await anyio.sleep(poll_interval_s)
--- a/src/exo/worker/utils/tests/test_macmon.py
+++ b/src/exo/worker/utils/tests/test_macmon.py
@@ -1,77 +0,0 @@
-"""Tests for macmon error handling.
-
-These tests verify that MacMon errors are handled gracefully without
-crashing the application or spamming logs.
-"""
-
-import platform
-from subprocess import CalledProcessError
-from unittest.mock import AsyncMock, patch
-
-import pytest
-
-from exo.worker.utils.macmon import MacMonError, get_metrics_async
-
-
-@pytest.mark.skipif(
-    platform.system().lower() != "darwin" or "arm" not in platform.machine().lower(),
-    reason="MacMon only supports macOS with Apple Silicon",
-)
-class TestMacMonErrorHandling:
-    """Test MacMon error handling."""
-
-    async def test_called_process_error_wrapped_as_macmon_error(self) -> None:
-        """CalledProcessError should be wrapped as MacMonError."""
-        mock_error = CalledProcessError(
-            returncode=1,
-            cmd=["macmon", "pipe", "-s", "1"],
-            stderr=b"some error message",
-        )
-
-        with (
-            patch(
-                "exo.worker.utils.macmon.shutil.which", return_value="/usr/bin/macmon"
-            ),
-            patch(
-                "exo.worker.utils.macmon.run_process", new_callable=AsyncMock
-            ) as mock_run,
-        ):
-            mock_run.side_effect = mock_error
-
-            with pytest.raises(MacMonError) as exc_info:
-                await get_metrics_async()
-
-            assert "MacMon failed with return code 1" in str(exc_info.value)
-            assert "some error message" in str(exc_info.value)
-
-    async def test_called_process_error_with_no_stderr(self) -> None:
-        """CalledProcessError with no stderr should be handled gracefully."""
-        mock_error = CalledProcessError(
-            returncode=1,
-            cmd=["macmon", "pipe", "-s", "1"],
-            stderr=None,
-        )
-
-        with (
-            patch(
-                "exo.worker.utils.macmon.shutil.which", return_value="/usr/bin/macmon"
-            ),
-            patch(
-                "exo.worker.utils.macmon.run_process", new_callable=AsyncMock
-            ) as mock_run,
-        ):
-            mock_run.side_effect = mock_error
-
-            with pytest.raises(MacMonError) as exc_info:
-                await get_metrics_async()
-
-            assert "MacMon failed with return code 1" in str(exc_info.value)
-            assert "no stderr" in str(exc_info.value)
-
-    async def test_macmon_not_found_raises_macmon_error(self) -> None:
-        """When macmon is not found in PATH, MacMonError should be raised."""
-        with patch("exo.worker.utils.macmon.shutil.which", return_value=None):
-            with pytest.raises(MacMonError) as exc_info:
-                await get_metrics_async()
-
-            assert "MacMon not found in PATH" in str(exc_info.value)
--- a/tests/headless_runner.py
+++ b/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,7 +34,8 @@ from exo.shared.types.worker.instances import (
 )
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
-from exo.utils.channels import MpReceiver, MpSender, mp_channel
+from exo.utils.channels import MpReceiver, MpSender, channel, mp_channel
+from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.worker.download.impl_shard_downloader import (
    build_full_shard,
    exo_shard_downloader,
@@ -65,6 +66,7 @@ async def main():
    app = FastAPI()
    app.post("/ring")(ring_backend)
    app.post("/jaccl")(jaccl_backend)
+    app.post("/tb_detection")(tb_detection)
    shutdown = anyio.Event()
    await serve(
        app,  # type: ignore
@@ -76,6 +78,15 @@ async def main():
    shutdown.set()


+async def tb_detection():
+    send, recv = channel[GatheredInfo]()
+    ig = InfoGatherer(send)
+    with anyio.move_on_after(1):
+        await ig._monitor_system_profiler()  # 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))
@@ -169,10 +180,16 @@ async def execute_test(test: Tests, instance: Instance, hn: str):
            send.send(StartWarmup(instance_id=iid))
            send.send(
                ChatCompletion(
-                    task_params=ResponsesRequest(
+                    task_params=ChatCompletionTaskParams(
                        model=test.model_id,
-                        instructions="You are a helpful assistant",
-                        input="What is the capital of France?",
+                        messages=[
+                            ChatCompletionMessage(
+                                role="system", content="You are a helpful assistant"
+                            ),
+                            ChatCompletionMessage(
+                                role="user", content="What is the capital of France?"
+                            ),
+                        ],
                    ),
                    command_id=CommandId("yo"),
                    instance_id=iid,
@@ -203,16 +220,16 @@ async def jaccl_backend(test: Tests):
            break
    else:
        raise ValueError(f"{weird_hn} not in {test.devs}")
-    return await execute_test(test, jaccl_instance(test, iid, hn), hn)
+    return await execute_test(test, jaccl_instance(test, iid), hn)


-def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
+def jaccl_instance(test: Tests, iid: InstanceId):
    meta = MODEL_CARDS[test.model_id].metadata
    world_size = len(test.devs)

    return MlxJacclInstance(
        instance_id=iid,
-        ibv_devices=[[None, "rdma_en3"], ["rdma_en3", None]],
+        jaccl_devices=[[None, "rdma_en3"], ["rdma_en3", None]],
        # rank 0 is always coordinator
        jaccl_coordinators={
            NodeId(host[0]): test.devs[0][1] + ":52416" for host in test.devs
Author	SHA1	Message	Date
Sami Khan	cd630dea43	formatting	2026-01-20 07:54:52 +05:00
Sami Khan	e55dae5ce8	code quality	2026-01-20 07:47:14 +05:00
Sami Khan	302c43afd5	Merge main into sami/flash	2026-01-20 07:31:22 +05:00
Sami Khan	2cf59e2322	nix flake check	2026-01-20 07:21:41 +05:00
Sami Khan	e506c7d65c	exo plugins	2026-01-20 06:53:43 +05:00
Sami Khan	c1fa2ddeaf	SLURM compatible commands	2026-01-20 06:53:43 +05:00
Alex Cheema	39f0ed6018	Prepend <think> tag to stream for thinking models like GLM-4.7 (#1186 ) ## Motivation For thinking models like GLM-4.7, the `<think>` tag is inserted by the tokenizer's `apply_chat_template()` into the prompt (input). The model generates tokens starting after this tag, so `<think>` never appears in the streamed output. The frontend expects `<think>...</think>` tags to extract and display thinking content. Log evidence: ``` [gMASK]<sop><\|system\|>...<\|user\|>...<\|assistant\|><think> ``` The prompt ends with `<think>`, so the model generates content after it, never returning the opening tag. ## Changes - Added `detect_thinking_prompt_suffix()` helper function in `utils_mlx.py` to detect if a prompt ends with `<think>` tag - Added `parse_thinking_models()` generator wrapper in `runner.py` that prepends the thinking tag to the output stream - Modified the main generation loop to use the thinking wrapper for non-GptOssModel models when a thinking prefix is detected - Updated test mocks to handle the new `apply_chat_template` call ## Why It Works The solution follows the same pattern as `parse_gpt_oss()` - a generator wrapper that transforms the output stream. When the chat template ends with `<think>`, we prepend this tag to the first generated token so the frontend receives the complete `<think>...</think>` structure it expects. ## Test Plan ### Manual Testing <!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB, connected via Thunderbolt 4) --> <!-- What you did: --> - Run exo: `uv run exo` - Send a chat request to GLM-4.7: ```bash curl http://localhost:52415/v1/chat/completions -H "Content-Type: application/json" -d '{ "model": "mlx-community/GLM-4.7-8bit-gs32", "messages": [{"role": "user", "content": "What is 2+2?"}], "stream": true }' ``` - Verify the streamed response starts with `<think>` tag - Verify the frontend dashboard correctly shows the thinking section collapsed ### Automated Testing - All 72 worker tests pass: `uv run pytest src/exo/worker/` - Type checker passes: `uv run basedpyright` - Linter passes: `uv run ruff check` 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com> Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>	2026-01-19 19:44:51 +00:00
Alex Cheema	ee43b598fe	Split NodePerformanceProfile into granular state mappings (#1209 ) ## Motivation The current `NodePerformanceProfile` is a monolithic object where every update (even 1-second memory updates) replaces the entire profile, touching unrelated data. Different fields update at vastly different frequencies: \| Data \| Update Frequency \| \|------\|------------------\| \| Memory, System \| 1 second \| \| Thunderbolt \| 5 seconds \| \| Network interfaces \| 10 seconds \| \| Friendly name \| 60 seconds \| \| Model/Chip ID \| Once at startup \| ## Changes Split into separate state mappings so each data type updates independently: - `node_identities`: Static and slow-changing data (model_id, chip_id, friendly_name) - `node_memory`: RAM and swap usage - `node_system`: GPU usage, temperature, power, CPU metrics - `node_network`: Network interface information - `node_thunderbolt`: Thunderbolt interface identifiers Added a backwards-compatible `node_profiles` property that reconstructs `NodePerformanceProfile` from the granular mappings for dashboard compatibility. Files modified: - `src/exo/shared/types/profiling.py` - Added `NodeIdentity`, `NodeNetworkInfo`, `NodeThunderboltInfo` types - `src/exo/shared/types/state.py` - Added 5 new mappings + `node_profiles` property - `src/exo/shared/apply.py` - Updated `apply_node_gathered_info` and `apply_node_timed_out` ## Why It Works Each info type now writes only to its specific mapping, avoiding unnecessary updates to unrelated data. The `MacThunderboltConnections` handler reads from `node_thunderbolt` instead of the old `node_profiles` for RDMA connection mapping. The backwards-compatible property ensures the dashboard continues to work unchanged. ## Test Plan ### Manual Testing <!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB, connected via Thunderbolt 4) --> <!-- What you did: --> - Start exo and verify dashboard shows node info - Verify memory/GPU updates stream correctly - Check that node timeout properly cleans up all mappings ### Automated Testing - All 162 existing tests pass - basedpyright: 0 errors - ruff check: All checks passed - nix fmt: Applied 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-19 18:24:15 +00:00
rltakashige	5fd55594c9	Wrap pipeline models for explicit mx.depends between cache and logits (#1206 ) ## Motivation GPU timeouts often when prompt size > profile_step_size. It also happens for seemingly random models. ## Changes Add mx.depends for cache on the logits. All gather at the model level rather than the layer level, reducing the amount of data sent. ## Why It Works mlx_lm's prefill loop only evaluates cache state, not logits. When prompt > prefill_step_size, the all_gather is never evaluated, causing GPU timeout. ## Test Plan ### Manual Testing <!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB, connected via Thunderbolt 4) --> <!-- What you did: --> <!-- - --> ### Automated Testing Added failing test cases and then resolved them.	2026-01-19 17:49:42 +00:00
Jake Hillion	5ab1f8b3e2	NetworkSetupHelper: detect stale startup script content The daemonAlreadyInstalled() function checked that the startup script file existed and validated plist properties, but did not compare the actual script content. If the setupScript constant was updated in a new app version, the stale on-disk script would not be detected or replaced. Added a guard clause that reads the installed script from disk and compares it against the expected setupScript content (with whitespace normalization). When content differs, the function returns false, triggering the reinstallation flow with an admin privileges prompt. Test plan: - Installed on a cluster that had the previous network config. Got the popup asking for permissions. After accepting I could run Kimi K2 Thinking Tensor RDMA on all 4 nodes.	2026-01-19 17:36:15 +00:00
Evan Quiney	2202685c3e	refactor all information sources (including ipless rdma discovery) (#928 ) ## Motivation Information gathering is tightly coupled to MacMon - we should start generalizing our information sources so we can add more in future. ## Changes Added a new system to gather any information. Currently, it is attached to the Worker - though this is mostly to keep the data processing logic simple. It could be made independent quite easily. I also refactored topology to include different kinds of connections as we can gather RDMA connections without having a pre-existing socket connection, and made the relevant placement updates. We should no longer need the network locations script in the app. Other sources of information now include: - static node information like "model" and "chip" (macos, "Unknown" fallback) - device friendly name (macos, falls back to device hostname) - network interfaces + ips (cross platform) - thunderbolt interfaces (macos) - thunderbolt connections (macos) - RAM usage (cross platform) - per-device configuration written to EXO_HOME/config.toml ## Limitations Model and Chip are not cross platform concepts. We do not differentiate between unified and non-unified memory systems. A lot of this data collection is based on simple timers. Watching the SC store on macos is the correct way to gather some of this information, but requires a detour into rust for macos. ## Why It Works The InfoGatherer is a generic subsystem which returns a union of metric datatypes. It writes them to an event, which is applied to state. It is currently re-spawned with the worker so each cluster receives the correct information. As for topology, macOS identifies TB ports with a uuid in SPThunderboltDataType, and also stores remote uuids if it can find them. These changes read that data with the system_profiler, hopefully not so often as to cause notable performance impacts (though this should be tuned) but frequently enough for moderate responsiveness. As we can identify TB connections between devices without needing ips attached to each interface, we can remove the network setup script (almost) completely. ## Test Plan ### Manual Testing Spawn RDMA instances without enabling DHCP on the RDMA interfaces. ### Automated Testing Updated the current master and shared tests to cover the topology refactor and new events. --------- Co-authored-by: Sami Khan <smsak99@gmail.com> Co-authored-by: Alex Cheema <alexcheema123@gmail.com> Co-authored-by: Jake Hillion <jake@hillion.co.uk>	2026-01-19 16:58:09 +00:00
Andrei Onel	ce3ad391b1	Update README.md with some changes from release 1.0.61 (#1157 ) Updated README.md with documentation for four new features: - added a "Benchmarking" section documenting the exo-bench tool for measuring model performance across different placement configurations - documented the custom namespace feature for cluster isolation in the macOS app section - added a "Configuration Options" subsection explaining the --no-worker CLI flag for coordinator-only nodes - added a "File Locations (Linux)" subsection documenting XDG Base Directory Specification compliance on Linux systems Issue #930	2026-01-19 16:43:18 +00:00
Jake Hillion	fb0151630d	shard_downloader: make on_progress callback async The on_progress callback was synchronous but always invoked from async contexts, forcing the use of send_nowait() which could raise WouldBlock if the channel buffer was full, potentially dropping progress updates. Changed the callback type from `Callable[[ShardMetadata, RepoDownloadProgress], None]` to return a coroutine, updated all implementations to be async, and replaced send_nowait() with await send() in the worker's download progress handler. This allows proper backpressure handling when sending download progress events through the channel, eliminating the "Footgun!" that was previously documented in the code. Test plan: - Built a DMG and ran it on one node. All existing models showed as downloaded. - Downloaded a new model. The progress bar on the download page worked. - Downloaded another new model. The progress bar on the home page worked.	2026-01-19 16:19:37 +00:00
Alex Cheema	346b13e2c9	Enhance LaTeX rendering in dashboard markdown (#1197 ) ## Motivation When models output LaTeX-formatted math proofs, the dashboard was not rendering them correctly. Issues included: - `\documentclass`, `\begin{document}`, `\usepackage` showing as raw text - `$...$` inline math with complex expressions (like `\frac`, `\ldots`) not rendering due to markdown escaping backslashes - `\begin{align}...\end{align}` and other math environments showing as raw text - `\emph{...}`, `\textbf{...}` LaTeX formatting commands not being converted - `$\require{...}$` (MathJax-specific) causing KaTeX errors - `\begin{proof}...\end{proof}` showing as raw text ## Changes Enhanced `MarkdownContent.svelte` with comprehensive LaTeX support: Math extraction before markdown processing: - Extract `$...$`, `$$...$$`, `\(...\)`, `\[...\]` into placeholders before markdown processes the text - Use alphanumeric placeholders (`MATHPLACEHOLDERINLINE0END`) that won't be interpreted as HTML tags - Restore and render with KaTeX after markdown processing LaTeX document command removal: - Strip `\documentclass{...}`, `\usepackage{...}`, `\begin{document}`, `\end{document}` - Strip `\maketitle`, `\title{...}`, `\author{...}`, `\date{...}` - Strip `\require{...}` (MathJax-specific, not KaTeX) - Replace `tikzpicture` environments with `[diagram]` placeholder - Strip `\label{...}` cross-reference commands LaTeX math environments: - Convert `\begin{align}`, `\begin{equation}`, `\begin{gather}`, etc. to display math blocks LaTeX text formatting:* - `\emph{...}` and `\textit{...}` → `<em>...</em>` - `\textbf{...}` → `<strong>...</strong>` - `\texttt{...}` → `<code>...</code>` - `\underline{...}` → `<u>...</u>` LaTeX environments styling: - `\begin{proof}...\end{proof}` → styled proof block with QED symbol - `\begin{theorem}`, `\begin{lemma}`, etc. → styled theorem blocks Display math enhancements: - Wrapped in styled container with subtle gold border - "LaTeX" label and copy button appear on hover - Dark theme KaTeX color overrides for better readability - Custom scrollbar for overflow ## Why It Works The key insight is that markdown processing was escaping backslashes in LaTeX before KaTeX could see them. By extracting all math expressions into alphanumeric placeholders before markdown runs, then restoring them after, the LaTeX content passes through to KaTeX unmodified. Using purely alphanumeric placeholders like `MATHPLACEHOLDERINLINE0END` instead of `<<MATH_INLINE_0>>` prevents markdown from interpreting them as HTML tags and stripping them. ## Test Plan ### Manual Testing - Hardware: Any machine with the dashboard - What you did: - Ask model to "write a proof in latex" - Verify inline math like `$x \in S$` renders correctly - Verify display math like `\begin{align}...\end{align}` renders as block - Verify `\documentclass`, `\begin{document}` are stripped (not shown) - Verify `\emph{...}` converts to italics - Verify copy button works on display math blocks - Test edge cases: `$5` (currency) stays as text, `\$50` (escaped) becomes `$50` Before: <img width="799" height="637" alt="Screenshot 2026-01-19 at 11 51 22 AM" src="https://github.com/user-attachments/assets/62a705b8-b3c2-47b8-afd0-5d0c1b240e44" /> After: <img width="809" height="642" alt="Screenshot 2026-01-19 at 11 46 58 AM" src="https://github.com/user-attachments/assets/4f35fa1d-333c-4285-bc68-58a50f8f148e" /> ### Automated Testing - Dashboard builds successfully with `npm run build` - Existing functionality preserved 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-19 14:50:41 +00:00
rltakashige	ea0588429b	Custom mlx layer composition (#1201 ) ## Motivation With a single pipeline layer, PipelineFirstLayer gets composed with PipelineLastLayer. ## Changes <!-- Describe what you changed in detail --> ## Why It Works <!-- Explain why your approach solves the problem --> ## Test Plan ### Manual Testing ### Automated Testing Made failing tests. Fixed them!	2026-01-19 12:36:25 +00:00
rltakashige	73b3f87e07	Set swa_idx and ga_idx for single layer (#1202 ) ## Motivation Layer types does not contain either "sliding_attention" or "full_attention" for pipeline parallel (single layer). ## Changes <!-- Describe what you changed in detail --> ## Why It Works <!-- Explain why your approach solves the problem --> ## Test Plan ### Manual Testing Manually tested single layer of GPT OSS. Doesn't crash ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> <!-- - -->	2026-01-19 12:31:11 +00:00
Sami Khan	37c5a2a246	Merge branch 'main' into sami/flash	2026-01-15 08:57:36 +05:00
Sami Khan	4d7f03834a	deleted separate server	2026-01-15 08:50:45 +05:00
Sami Khan	bdb9fbc8c0	Merge branch 'main' into sami/flash	2026-01-14 08:10:51 +05:00
Sami Khan	8c7180810c	type checking	2026-01-14 07:15:45 +05:00
Sami Khan	318c6e000b	code cleanup	2026-01-14 04:56:59 +05:00
Sami Khan	2d45544da0	use rsh server instead of ssh	2026-01-13 02:46:25 +05:00
Sami Khan	7cbafa768a	flash+exo	2026-01-12 10:26:16 +05:00
				`@@ -1 +0,0 @@`
				`"""API adapters for different API formats (Claude, OpenAI Responses, etc.)."""`