From 19965c7ba54dcab1f6446e01e87f9a8610dcce7c Mon Sep 17 00:00:00 2001 From: ciaranbor <81697641+ciaranbor@users.noreply.github.com> Date: Mon, 2 Feb 2026 17:19:45 +0000 Subject: [PATCH] Ciaran/profiling (#1345) ## Motivation Know what the hell is going on ## Changes - Tracing library (src/exo/shared/tracing.py): trace() context manager, Chrome Trace Format export, statistics computation - Runner instrumentation (src/exo/worker/engines/image/pipeline/runner.py): Wrapped sync/async steps, compute blocks, and send/recv operations - Trace collection: Workers send traces to master after task completion; merged into ~/.exo/traces/trace_{task_id}.json - API endpoints: List, fetch, stats, and raw download at /v1/traces/* - Dashboard: Trace list and detail pages with Perfetto integration ## Why It Works Screenshot 2026-01-30 at 19 00 09 Screenshot 2026-01-30 at 19 00 58 --- dashboard/src/lib/stores/app.svelte.ts | 82 ++++ dashboard/src/routes/traces/+page.svelte | 190 +++++++++ .../src/routes/traces/[taskId]/+page.svelte | 367 ++++++++++++++++++ src/exo/master/api.py | 134 +++++++ src/exo/master/main.py | 64 ++- src/exo/shared/apply.py | 9 +- src/exo/shared/constants.py | 3 + src/exo/shared/tracing.py | 238 ++++++++++++ src/exo/shared/types/api.py | 42 ++ src/exo/shared/types/events.py | 27 +- .../worker/engines/image/pipeline/runner.py | 342 ++++++++++------ src/exo/worker/runner/runner.py | 43 +- 12 files changed, 1407 insertions(+), 134 deletions(-) create mode 100644 dashboard/src/routes/traces/+page.svelte create mode 100644 dashboard/src/routes/traces/[taskId]/+page.svelte create mode 100644 src/exo/shared/tracing.py diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 9fa371fb..51de6c66 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -178,6 +178,36 @@ interface ImageApiResponse { data: Array<{ b64_json?: string; url?: string }>; } +// Trace API response types +export interface TraceCategoryStats { + totalUs: number; + count: number; + minUs: number; + maxUs: number; + avgUs: number; +} + +export interface TraceRankStats { + byCategory: Record; +} + +export interface TraceStatsResponse { + taskId: string; + totalWallTimeUs: number; + byCategory: Record; + byRank: Record; +} + +export interface TraceListItem { + taskId: string; + createdAt: string; + fileSize: number; +} + +export interface TraceListResponse { + traces: TraceListItem[]; +} + interface RawStateResponse { topology?: RawTopology; instances?: Record< @@ -2555,6 +2585,49 @@ class AppStore { throw error; } } + + /** + * List all available traces + */ + async listTraces(): Promise { + const response = await fetch("/v1/traces"); + if (!response.ok) { + throw new Error(`Failed to list traces: ${response.status}`); + } + return (await response.json()) as TraceListResponse; + } + + /** + * Check if a trace exists for a given task ID + */ + async checkTraceExists(taskId: string): Promise { + try { + const response = await fetch(`/v1/traces/${encodeURIComponent(taskId)}`); + return response.ok; + } catch { + return false; + } + } + + /** + * Get computed statistics for a task's trace + */ + async fetchTraceStats(taskId: string): Promise { + const response = await fetch( + `/v1/traces/${encodeURIComponent(taskId)}/stats`, + ); + if (!response.ok) { + throw new Error(`Failed to fetch trace stats: ${response.status}`); + } + return (await response.json()) as TraceStatsResponse; + } + + /** + * Get the URL for the raw trace file (for Perfetto) + */ + getTraceRawUrl(taskId: string): string { + return `/v1/traces/${encodeURIComponent(taskId)}/raw`; + } } export const appStore = new AppStore(); @@ -2666,3 +2739,12 @@ export const startDownload = (nodeId: string, shardMetadata: object) => appStore.startDownload(nodeId, shardMetadata); export const deleteDownload = (nodeId: string, modelId: string) => appStore.deleteDownload(nodeId, modelId); + +// Trace actions +export const listTraces = () => appStore.listTraces(); +export const checkTraceExists = (taskId: string) => + appStore.checkTraceExists(taskId); +export const fetchTraceStats = (taskId: string) => + appStore.fetchTraceStats(taskId); +export const getTraceRawUrl = (taskId: string) => + appStore.getTraceRawUrl(taskId); diff --git a/dashboard/src/routes/traces/+page.svelte b/dashboard/src/routes/traces/+page.svelte new file mode 100644 index 00000000..dae70ff5 --- /dev/null +++ b/dashboard/src/routes/traces/+page.svelte @@ -0,0 +1,190 @@ + + +
+ +
+
+
+

+ Traces +

+
+
+ +
+
+ + {#if loading} +
+
Loading traces...
+
+ {:else if error} +
+
{error}
+
+ {:else if traces.length === 0} +
+
No traces found.
+
+ Run exo with EXO_TRACING_ENABLED=1 to collect traces. +
+
+ {:else} +
+ {#each traces as trace} +
+
+ + {trace.taskId} + +
+ {formatDate(trace.createdAt)} • {formatBytes( + trace.fileSize, + )} +
+
+
+ + View Stats + + + +
+
+ {/each} +
+ {/if} +
+
diff --git a/dashboard/src/routes/traces/[taskId]/+page.svelte b/dashboard/src/routes/traces/[taskId]/+page.svelte new file mode 100644 index 00000000..f370c8fa --- /dev/null +++ b/dashboard/src/routes/traces/[taskId]/+page.svelte @@ -0,0 +1,367 @@ + + +
+ +
+
+
+

+ Trace +

+

+ {taskId} +

+
+
+ + All Traces + + + +
+
+ + {#if loading} +
+
Loading trace data...
+
+ {:else if error} +
+
{error}
+
+ {:else if stats} + +
+

+ Summary +

+
+ {formatDuration(stats.totalWallTimeUs)} +
+
Total wall time
+
+ + + {#if phases.length > 0} +
+

+ By Phase (avg per node) +

+
+ {#each phases as phase} + {@const normalizedTotal = phase.totalUs / nodeCount} + {@const normalizedStepCount = phase.stepCount / nodeCount} +
+
+ {phase.name} + + {formatDuration(normalizedTotal)} + + ({normalizedStepCount} steps, {formatDuration( + normalizedTotal / normalizedStepCount, + )}/step) + + +
+ {#if phase.subcategories.length > 0} +
+ {#each phase.subcategories as subcat} + {@const normalizedSubcat = + subcat.stats.totalUs / nodeCount} + {@const pct = formatPercentage( + normalizedSubcat, + normalizedTotal, + )} + {@const perStep = normalizedSubcat / normalizedStepCount} +
+ {subcat.name} + + {formatDuration(normalizedSubcat)} + ({pct}) + {formatDuration(perStep)}/step + +
+ +
+
+
+ {/each} +
+ {/if} +
+ {/each} +
+
+ {/if} + + + {#if sortedRanks.length > 0} +
+

+ By Rank +

+
+ {#each sortedRanks as rank} + {@const rankStats = stats.byRank[rank]} + {@const rankPhases = parsePhases(rankStats.byCategory)} +
+
+ Rank {rank} +
+
+ {#each rankPhases as phase} +
+
+ {phase.name} + + {formatDuration(phase.totalUs)} + + ({phase.stepCount}x) + + +
+ {#if phase.subcategories.length > 0} +
+ {#each phase.subcategories as subcat} + {@const pct = formatPercentage( + subcat.stats.totalUs, + phase.totalUs, + )} + {@const perStep = + subcat.stats.totalUs / phase.stepCount} +
+ {subcat.name} + + {formatDuration(subcat.stats.totalUs)} + ({pct}) + {formatDuration(perStep)}/step + +
+ {/each} +
+ {/if} +
+ {/each} +
+
+ {/each} +
+
+ {/if} + {/if} +
+
diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 908716d8..4855e62a 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -3,7 +3,9 @@ import contextlib import json import time from collections.abc import AsyncGenerator, Awaitable, Callable +from datetime import datetime, timezone from http import HTTPStatus +from pathlib import Path from typing import Annotated, Literal, cast from uuid import uuid4 @@ -40,6 +42,7 @@ from exo.shared.apply import apply from exo.shared.constants import ( EXO_IMAGE_CACHE_DIR, EXO_MAX_CHUNK_SIZE, + EXO_TRACING_CACHE_DIR, ) from exo.shared.election import ElectionMessage from exo.shared.logging import InterceptLogger @@ -48,6 +51,7 @@ from exo.shared.models.model_cards import ( ModelCard, ModelId, ) +from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file from exo.shared.types.api import ( AdvancedImageParams, BenchChatCompletionRequest, @@ -81,6 +85,13 @@ from exo.shared.types.api import ( StartDownloadParams, StartDownloadResponse, ToolCall, + TraceCategoryStats, + TraceEventResponse, + TraceListItem, + TraceListResponse, + TraceRankStats, + TraceResponse, + TraceStatsResponse, ) from exo.shared.types.chunks import ( ErrorChunk, @@ -115,6 +126,7 @@ from exo.shared.types.events import ( Event, ForwarderEvent, IndexedEvent, + TracesMerged, ) from exo.shared.types.memory import Memory from exo.shared.types.openai_responses import ( @@ -275,6 +287,10 @@ class API: self.app.get("/events")(lambda: self._event_log) self.app.post("/download/start")(self.start_download) self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download) + self.app.get("/v1/traces")(self.list_traces) + self.app.get("/v1/traces/{task_id}")(self.get_trace) + self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats) + self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw) async def place_instance(self, payload: PlaceInstanceParams): command = PlaceInstance( @@ -1279,6 +1295,24 @@ class API: except BrokenResourceError: self._text_generation_queues.pop(event.command_id, None) + if isinstance(event, TracesMerged): + self._save_merged_trace(event) + + def _save_merged_trace(self, event: TracesMerged) -> None: + traces = [ + TraceEvent( + name=t.name, + start_us=t.start_us, + duration_us=t.duration_us, + rank=t.rank, + category=t.category, + ) + for t in event.traces + ] + output_path = EXO_TRACING_CACHE_DIR / f"trace_{event.task_id}.json" + export_trace(traces, output_path) + logger.debug(f"Saved merged trace to {output_path}") + async def _pause_on_new_election(self): with self.election_receiver as ems: async for message in ems: @@ -1325,3 +1359,103 @@ class API: ) await self._send_download(command) return DeleteDownloadResponse(command_id=command.command_id) + + def _get_trace_path(self, task_id: str) -> Path: + return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json" + + async def list_traces(self) -> TraceListResponse: + traces: list[TraceListItem] = [] + + for trace_file in sorted( + EXO_TRACING_CACHE_DIR.glob("trace_*.json"), + key=lambda p: p.stat().st_mtime, + reverse=True, + ): + # Extract task_id from filename (trace_{task_id}.json) + task_id = trace_file.stem.removeprefix("trace_") + stat = trace_file.stat() + created_at = datetime.fromtimestamp( + stat.st_mtime, tz=timezone.utc + ).isoformat() + traces.append( + TraceListItem( + task_id=task_id, + created_at=created_at, + file_size=stat.st_size, + ) + ) + + return TraceListResponse(traces=traces) + + async def get_trace(self, task_id: str) -> TraceResponse: + trace_path = self._get_trace_path(task_id) + + if not trace_path.exists(): + raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}") + + trace_events = load_trace_file(trace_path) + + return TraceResponse( + task_id=task_id, + traces=[ + TraceEventResponse( + name=event.name, + start_us=event.start_us, + duration_us=event.duration_us, + rank=event.rank, + category=event.category, + ) + for event in trace_events + ], + ) + + async def get_trace_stats(self, task_id: str) -> TraceStatsResponse: + trace_path = self._get_trace_path(task_id) + + if not trace_path.exists(): + raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}") + + trace_events = load_trace_file(trace_path) + stats = compute_stats(trace_events) + + return TraceStatsResponse( + task_id=task_id, + total_wall_time_us=stats.total_wall_time_us, + by_category={ + category: TraceCategoryStats( + total_us=cat_stats.total_us, + count=cat_stats.count, + min_us=cat_stats.min_us, + max_us=cat_stats.max_us, + avg_us=cat_stats.avg_us, + ) + for category, cat_stats in stats.by_category.items() + }, + by_rank={ + rank: TraceRankStats( + by_category={ + category: TraceCategoryStats( + total_us=cat_stats.total_us, + count=cat_stats.count, + min_us=cat_stats.min_us, + max_us=cat_stats.max_us, + avg_us=cat_stats.avg_us, + ) + for category, cat_stats in rank_stats.items() + } + ) + for rank, rank_stats in stats.by_rank.items() + }, + ) + + async def get_trace_raw(self, task_id: str) -> FileResponse: + trace_path = self._get_trace_path(task_id) + + if not trace_path.exists(): + raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}") + + return FileResponse( + path=trace_path, + media_type="application/json", + filename=f"trace_{task_id}.json", + ) diff --git a/src/exo/master/main.py b/src/exo/master/main.py index d99a96c9..a3041143 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -11,6 +11,7 @@ from exo.master.placement import ( place_instance, ) from exo.shared.apply import apply +from exo.shared.constants import EXO_TRACING_ENABLED from exo.shared.types.commands import ( CreateInstance, DeleteInstance, @@ -35,6 +36,9 @@ from exo.shared.types.events import ( NodeTimedOut, TaskCreated, TaskDeleted, + TraceEventData, + TracesCollected, + TracesMerged, ) from exo.shared.types.state import State from exo.shared.types.tasks import ( @@ -86,6 +90,8 @@ class Master: self._multi_buffer = MultiSourceBuffer[NodeId, Event]() # TODO: not have this self._event_log: list[Event] = [] + self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} + self._expected_ranks: dict[TaskId, set[int]] = {} async def run(self): logger.info("Starting Master") @@ -187,13 +193,14 @@ class Master: ) task_id = TaskId() + selected_instance_id = available_instance_ids[0] generated_events.append( TaskCreated( task_id=task_id, task=ImageGenerationTask( task_id=task_id, command_id=command.command_id, - instance_id=available_instance_ids[0], + instance_id=selected_instance_id, task_status=TaskStatus.Pending, task_params=command.task_params, ), @@ -201,6 +208,17 @@ class Master: ) self.command_task_mapping[command.command_id] = task_id + + if EXO_TRACING_ENABLED: + selected_instance = self.state.instances.get( + selected_instance_id + ) + if selected_instance: + ranks = set( + shard.device_rank + for shard in selected_instance.shard_assignments.runner_to_shard.values() + ) + self._expected_ranks[task_id] = ranks case ImageEdits(): for instance in self.state.instances.values(): if ( @@ -229,13 +247,14 @@ class Master: ) task_id = TaskId() + selected_instance_id = available_instance_ids[0] generated_events.append( TaskCreated( task_id=task_id, task=ImageEditsTask( task_id=task_id, command_id=command.command_id, - instance_id=available_instance_ids[0], + instance_id=selected_instance_id, task_status=TaskStatus.Pending, task_params=command.task_params, ), @@ -243,6 +262,17 @@ class Master: ) self.command_task_mapping[command.command_id] = task_id + + if EXO_TRACING_ENABLED: + selected_instance = self.state.instances.get( + selected_instance_id + ) + if selected_instance: + ranks = set( + shard.device_rank + for shard in selected_instance.shard_assignments.runner_to_shard.values() + ) + self._expected_ranks[task_id] = ranks case DeleteInstance(): placement = delete_instance(command, self.state.instances) transition_events = get_transition_events( @@ -335,6 +365,10 @@ class Master: local_event.origin, ) for event in self._multi_buffer.drain(): + if isinstance(event, TracesCollected): + await self._handle_traces_collected(event) + continue + logger.debug(f"Master indexing event: {str(event)[:100]}") indexed = IndexedEvent(event=event, idx=len(self._event_log)) self.state = apply(self.state, indexed) @@ -373,3 +407,29 @@ class Master: event=event.event, ) ) + + async def _handle_traces_collected(self, event: TracesCollected) -> None: + task_id = event.task_id + if task_id not in self._pending_traces: + self._pending_traces[task_id] = {} + self._pending_traces[task_id][event.rank] = event.traces + + if ( + task_id in self._expected_ranks + and set(self._pending_traces[task_id].keys()) + >= self._expected_ranks[task_id] + ): + await self._merge_and_save_traces(task_id) + + async def _merge_and_save_traces(self, task_id: TaskId) -> None: + all_trace_data: list[TraceEventData] = [] + for trace_data in self._pending_traces[task_id].values(): + all_trace_data.extend(trace_data) + + await self.event_sender.send( + TracesMerged(task_id=task_id, traces=all_trace_data) + ) + + del self._pending_traces[task_id] + if task_id in self._expected_ranks: + del self._expected_ranks[task_id] diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 675b9647..8c2304fe 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -25,6 +25,8 @@ from exo.shared.types.events import ( TestEvent, TopologyEdgeCreated, TopologyEdgeDeleted, + TracesCollected, + TracesMerged, ) from exo.shared.types.profiling import ( NodeIdentity, @@ -55,7 +57,12 @@ def event_apply(event: Event, state: State) -> State: """Apply an event to state.""" match event: case ( - TestEvent() | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived() + TestEvent() + | ChunkGenerated() + | TaskAcknowledged() + | InputChunkReceived() + | TracesCollected() + | TracesMerged() ): # Pass-through events that don't modify state return state case InstanceCreated(): diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py index dfd610b4..ad503e90 100644 --- a/src/exo/shared/constants.py +++ b/src/exo/shared/constants.py @@ -49,7 +49,10 @@ LIBP2P_COMMANDS_TOPIC = "commands" EXO_MAX_CHUNK_SIZE = 512 * 1024 EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images" +EXO_TRACING_CACHE_DIR = EXO_CACHE_HOME / "traces" EXO_ENABLE_IMAGE_MODELS = ( os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true" ) + +EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true" diff --git a/src/exo/shared/tracing.py b/src/exo/shared/tracing.py new file mode 100644 index 00000000..a353d923 --- /dev/null +++ b/src/exo/shared/tracing.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import json +import time +from collections import defaultdict +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field +from pathlib import Path +from typing import cast, final + +from exo.shared.constants import EXO_TRACING_ENABLED +from exo.worker.runner.bootstrap import logger + +# Context variable to track the current trace category for hierarchical nesting +_current_category: ContextVar[str | None] = ContextVar("current_category", default=None) + + +@final +@dataclass(frozen=True) +class TraceEvent: + name: str + start_us: int + duration_us: int + rank: int + category: str + + +@final +@dataclass +class CategoryStats: + total_us: int = 0 + count: int = 0 + min_us: int = 0 + max_us: int = 0 + + def add(self, duration_us: int) -> None: + if self.count == 0: + self.min_us = duration_us + self.max_us = duration_us + else: + self.min_us = min(self.min_us, duration_us) + self.max_us = max(self.max_us, duration_us) + self.total_us += duration_us + self.count += 1 + + @property + def avg_us(self) -> float: + return self.total_us / self.count if self.count > 0 else 0.0 + + +@final +@dataclass +class TraceStats: + total_wall_time_us: int = 0 + by_category: dict[str, CategoryStats] = field(default_factory=dict) + by_rank: dict[int, dict[str, CategoryStats]] = field(default_factory=dict) + + +# Global trace buffer - each rank accumulates traces here +_trace_buffer: list[TraceEvent] = [] + + +def _record_span( + name: str, start_us: int, duration_us: int, rank: int, category: str +) -> None: + _trace_buffer.append( + TraceEvent( + name=name, + start_us=start_us, + duration_us=duration_us, + rank=rank, + category=category, + ) + ) + + +@contextmanager +def trace( + name: str, + rank: int, + category: str = "compute", +) -> Generator[None, None, None]: + """Context manager to trace any operation. + + Nested traces automatically inherit the parent category, creating hierarchical + categories like "sync/compute" or "async/comms". + + Args: + name: Name of the operation (e.g., "recv 0", "send 1", "joint_blocks") + rank: This rank's ID + category: Category for grouping in trace viewer ("comm", "compute", "step") + + Example: + with trace(f"sync {t}", rank, "sync"): + with trace("joint_blocks", rank, "compute"): + # Recorded with category "sync/compute" + hidden_states = some_computation(...) + """ + if not EXO_TRACING_ENABLED: + yield + return + + # Combine with parent category if nested + parent = _current_category.get() + full_category = f"{parent}/{category}" if parent else category + + # Set as current for nested traces + token = _current_category.set(full_category) + + try: + start_us = int(time.time() * 1_000_000) + start_perf = time.perf_counter() + yield + duration_us = int((time.perf_counter() - start_perf) * 1_000_000) + _record_span(name, start_us, duration_us, rank, full_category) + finally: + _current_category.reset(token) + + +def get_trace_buffer() -> list[TraceEvent]: + return list(_trace_buffer) + + +def clear_trace_buffer() -> None: + _trace_buffer.clear() + + +def export_trace(traces: list[TraceEvent], output_path: Path) -> None: + trace_events: list[dict[str, object]] = [] + + for event in traces: + # Chrome trace format uses "X" for complete events (with duration) + chrome_event: dict[str, object] = { + "name": event.name, + "cat": event.category, + "ph": "X", + "ts": event.start_us, + "dur": event.duration_us, + "pid": 0, + "tid": event.rank, + "args": {"rank": event.rank}, + } + trace_events.append(chrome_event) + + ranks_seen = set(t.rank for t in traces) + for rank in ranks_seen: + trace_events.append( + { + "name": "thread_name", + "ph": "M", # Metadata event + "pid": 0, + "tid": rank, + "args": {"name": f"Rank {rank}"}, + } + ) + + chrome_trace = {"traceEvents": trace_events} + + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(chrome_trace, f, indent=2) + except OSError as e: + logger.warning("Failed to export trace to %s: %s", output_path, e) + + +def load_trace_file(path: Path) -> list[TraceEvent]: + with open(path) as f: + data = cast(dict[str, list[dict[str, object]]], json.load(f)) + + events = data.get("traceEvents", []) + traces: list[TraceEvent] = [] + + for event in events: + # Skip metadata events + if event.get("ph") == "M": + continue + + name = str(event.get("name", "")) + category = str(event.get("cat", "")) + ts_value = event.get("ts", 0) + dur_value = event.get("dur", 0) + tid_value = event.get("tid", 0) + start_us = int(ts_value) if isinstance(ts_value, (int, float, str)) else 0 + duration_us = int(dur_value) if isinstance(dur_value, (int, float, str)) else 0 + + # Get rank from tid or args + rank = int(tid_value) if isinstance(tid_value, (int, float, str)) else 0 + args = event.get("args") + if isinstance(args, dict): + args_dict = cast(dict[str, object], args) + rank_from_args = args_dict.get("rank") + if isinstance(rank_from_args, (int, float, str)): + rank = int(rank_from_args) + + traces.append( + TraceEvent( + name=name, + start_us=start_us, + duration_us=duration_us, + rank=rank, + category=category, + ) + ) + + return traces + + +def compute_stats(traces: list[TraceEvent]) -> TraceStats: + stats = TraceStats() + + if not traces: + return stats + + # Calculate wall time from earliest start to latest end + min_start = min(t.start_us for t in traces) + max_end = max(t.start_us + t.duration_us for t in traces) + stats.total_wall_time_us = max_end - min_start + + # Initialize nested dicts + by_category: dict[str, CategoryStats] = defaultdict(CategoryStats) + by_rank: dict[int, dict[str, CategoryStats]] = defaultdict( + lambda: defaultdict(CategoryStats) + ) + + for event in traces: + # By category + by_category[event.category].add(event.duration_us) + + # By rank and category + by_rank[event.rank][event.category].add(event.duration_us) + + stats.by_category = dict(by_category) + stats.by_rank = {k: dict(v) for k, v in by_rank.items()} + + return stats diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 3f665179..40dbb288 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -350,3 +350,45 @@ class StartDownloadResponse(CamelCaseModel): class DeleteDownloadResponse(CamelCaseModel): command_id: CommandId + + +class TraceEventResponse(CamelCaseModel): + name: str + start_us: int + duration_us: int + rank: int + category: str + + +class TraceResponse(CamelCaseModel): + task_id: str + traces: list[TraceEventResponse] + + +class TraceCategoryStats(CamelCaseModel): + total_us: int + count: int + min_us: int + max_us: int + avg_us: float + + +class TraceRankStats(CamelCaseModel): + by_category: dict[str, TraceCategoryStats] + + +class TraceStatsResponse(CamelCaseModel): + task_id: str + total_wall_time_us: int + by_category: dict[str, TraceCategoryStats] + by_rank: dict[int, TraceRankStats] + + +class TraceListItem(CamelCaseModel): + task_id: str + created_at: str + file_size: int + + +class TraceListResponse(CamelCaseModel): + traces: list[TraceListItem] diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 647e510f..5cf93d0c 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -1,4 +1,5 @@ from datetime import datetime +from typing import final from pydantic import Field @@ -10,7 +11,7 @@ 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 +from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel class EventId(Id): @@ -109,6 +110,28 @@ class TopologyEdgeDeleted(BaseEvent): conn: Connection +@final +class TraceEventData(FrozenModel): + name: str + start_us: int + duration_us: int + rank: int + category: str + + +@final +class TracesCollected(BaseEvent): + task_id: TaskId + rank: int + traces: list[TraceEventData] + + +@final +class TracesMerged(BaseEvent): + task_id: TaskId + traces: list[TraceEventData] + + Event = ( TestEvent | TaskCreated @@ -127,6 +150,8 @@ Event = ( | InputChunkReceived | TopologyEdgeCreated | TopologyEdgeDeleted + | TracesCollected + | TracesMerged ) diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py index 698887b0..e1f65efd 100644 --- a/src/exo/worker/engines/image/pipeline/runner.py +++ b/src/exo/worker/engines/image/pipeline/runner.py @@ -6,6 +6,11 @@ from mflux.models.common.config.config import Config from mflux.utils.exceptions import StopImageGenerationException from tqdm import tqdm +from exo.shared.constants import EXO_TRACING_ENABLED +from exo.shared.tracing import ( + clear_trace_buffer, + trace, +) from exo.shared.types.worker.shards import PipelineShardMetadata from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ( @@ -324,6 +329,7 @@ class DiffusionRunner: capture_steps = set() self._reset_all_caches() + clear_trace_buffer() time_steps = tqdm(range(runtime_config.num_inference_steps)) @@ -465,20 +471,22 @@ class DiffusionRunner: if self.group is None: return self._single_node_step(t, config, latents, prompt_data) elif t < config.init_time_step + num_sync_steps: - return self._sync_pipeline_step( - t, - config, - latents, - prompt_data, - ) + with trace(name=f"sync {t}", rank=self.rank, category="sync"): + return self._sync_pipeline_step( + t, + config, + latents, + prompt_data, + ) else: - return self._async_pipeline_step( - t, - config, - latents, - prompt_data, - is_first_async_step=t == config.init_time_step + num_sync_steps, - ) + with trace(name=f"async {t}", rank=self.rank, category="async"): + return self._async_pipeline_step( + t, + config, + latents, + prompt_data, + is_first_async_step=t == config.init_time_step + num_sync_steps, + ) def _single_node_step( self, @@ -586,30 +594,41 @@ class DiffusionRunner: if self.has_joint_blocks: if not self.is_first_stage: - hidden_states = mx.distributed.recv( - (batch_size, num_img_tokens, hidden_dim), - dtype, - self.prev_rank, - group=self.group, - ) - encoder_hidden_states = mx.distributed.recv( - (batch_size, text_seq_len, hidden_dim), - dtype, - self.prev_rank, - group=self.group, - ) - mx.eval(hidden_states, encoder_hidden_states) + with trace( + name=f"recv {self.prev_rank}", rank=self.rank, category="comms" + ): + hidden_states = mx.distributed.recv( + (batch_size, num_img_tokens, hidden_dim), + dtype, + self.prev_rank, + group=self.group, + ) + encoder_hidden_states = mx.distributed.recv( + (batch_size, text_seq_len, hidden_dim), + dtype, + self.prev_rank, + group=self.group, + ) + mx.eval(hidden_states, encoder_hidden_states) assert self.joint_block_wrappers is not None assert encoder_hidden_states is not None - for wrapper in self.joint_block_wrappers: - wrapper.set_patch(BlockWrapperMode.CACHING) - encoder_hidden_states, hidden_states = wrapper( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_embeddings, - ) + with trace( + name="joint_blocks", + rank=self.rank, + category="compute", + ): + for wrapper in self.joint_block_wrappers: + wrapper.set_patch(BlockWrapperMode.CACHING) + encoder_hidden_states, hidden_states = wrapper( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + text_embeddings=text_embeddings, + rotary_embeddings=image_rotary_embeddings, + ) + + if EXO_TRACING_ENABLED: + mx.eval(encoder_hidden_states, hidden_states) if self.owns_concat_stage: assert encoder_hidden_states is not None @@ -620,45 +639,63 @@ class DiffusionRunner: if self.has_single_blocks or self.is_last_stage: hidden_states = concatenated else: - concatenated = mx.distributed.send( - concatenated, self.next_rank, group=self.group - ) - mx.async_eval(concatenated) + with trace( + name=f"send {self.next_rank}", rank=self.rank, category="comms" + ): + concatenated = mx.distributed.send( + concatenated, self.next_rank, group=self.group + ) + mx.async_eval(concatenated) elif self.has_joint_blocks and not self.is_last_stage: assert encoder_hidden_states is not None - hidden_states = mx.distributed.send( - hidden_states, self.next_rank, group=self.group - ) - encoder_hidden_states = mx.distributed.send( - encoder_hidden_states, self.next_rank, group=self.group - ) - mx.async_eval(hidden_states, encoder_hidden_states) - - if self.has_single_blocks: - if not self.owns_concat_stage and not self.is_first_stage: - hidden_states = mx.distributed.recv( - (batch_size, text_seq_len + num_img_tokens, hidden_dim), - dtype, - self.prev_rank, - group=self.group, - ) - mx.eval(hidden_states) - - assert self.single_block_wrappers is not None - for wrapper in self.single_block_wrappers: - wrapper.set_patch(BlockWrapperMode.CACHING) - hidden_states = wrapper( - hidden_states=hidden_states, - text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_embeddings, - ) - - if not self.is_last_stage: + with trace(name=f"send {self.next_rank}", rank=self.rank, category="comms"): hidden_states = mx.distributed.send( hidden_states, self.next_rank, group=self.group ) - mx.async_eval(hidden_states) + encoder_hidden_states = mx.distributed.send( + encoder_hidden_states, self.next_rank, group=self.group + ) + mx.async_eval(hidden_states, encoder_hidden_states) + + if self.has_single_blocks: + if not self.owns_concat_stage and not self.is_first_stage: + with trace( + name=f"recv {self.prev_rank}", rank=self.rank, category="comms" + ): + hidden_states = mx.distributed.recv( + (batch_size, text_seq_len + num_img_tokens, hidden_dim), + dtype, + self.prev_rank, + group=self.group, + ) + mx.eval(hidden_states) + + assert self.single_block_wrappers is not None + with trace( + name="single blocks", + rank=self.rank, + category="compute", + ): + for wrapper in self.single_block_wrappers: + wrapper.set_patch(BlockWrapperMode.CACHING) + hidden_states = wrapper( + hidden_states=hidden_states, + text_embeddings=text_embeddings, + rotary_embeddings=image_rotary_embeddings, + ) + + if EXO_TRACING_ENABLED: + mx.eval(hidden_states) + + if not self.is_last_stage: + with trace( + name=f"send {self.next_rank}", rank=self.rank, category="comms" + ): + hidden_states = mx.distributed.send( + hidden_states, self.next_rank, group=self.group + ) + mx.async_eval(hidden_states) hidden_states = hidden_states[:, text_seq_len:, ...] @@ -742,14 +779,20 @@ class DiffusionRunner: ) if not self.is_first_stage: - hidden_states = mx.distributed.send(hidden_states, 0, group=self.group) - mx.async_eval(hidden_states) + with trace(name="send 0", rank=self.rank, category="comms"): + hidden_states = mx.distributed.send( + hidden_states, 0, group=self.group + ) + mx.async_eval(hidden_states) elif self.is_first_stage: - hidden_states = mx.distributed.recv_like( - prev_latents, src=self.world_size - 1, group=self.group - ) - mx.eval(hidden_states) + with trace( + name=f"recv {self.world_size - 1}", rank=self.rank, category="comms" + ): + hidden_states = mx.distributed.recv_like( + prev_latents, src=self.world_size - 1, group=self.group + ) + mx.eval(hidden_states) else: hidden_states = prev_latents @@ -809,10 +852,13 @@ class DiffusionRunner: and not self.is_last_stage and not is_first_async_step ): - patch = mx.distributed.recv_like( - patch, src=self.prev_rank, group=self.group - ) - mx.eval(patch) + with trace( + name=f"recv {self.prev_rank}", rank=self.rank, category="comms" + ): + patch = mx.distributed.recv_like( + patch, src=self.prev_rank, group=self.group + ) + mx.eval(patch) step_patch = mx.concatenate([patch, patch], axis=0) if needs_cfg else patch @@ -843,10 +889,13 @@ class DiffusionRunner: ) if not self.is_first_stage and t != config.num_inference_steps - 1: - patch_latents[patch_idx] = mx.distributed.send( - patch_latents[patch_idx], self.next_rank, group=self.group - ) - mx.async_eval(patch_latents[patch_idx]) + with trace( + name=f"send {self.next_rank}", rank=self.rank, category="comms" + ): + patch_latents[patch_idx] = mx.distributed.send( + patch_latents[patch_idx], self.next_rank, group=self.group + ) + mx.async_eval(patch_latents[patch_idx]) return mx.concatenate(patch_latents, axis=1) @@ -885,22 +934,28 @@ class DiffusionRunner: if self.has_joint_blocks: if not self.is_first_stage: patch_len = patch.shape[1] - patch = mx.distributed.recv( - (batch_size, patch_len, hidden_dim), - patch.dtype, - self.prev_rank, - group=self.group, - ) - mx.eval(patch) - - if patch_idx == 0: - encoder_hidden_states = mx.distributed.recv( - (batch_size, text_seq_len, hidden_dim), + with trace( + name=f"recv {self.prev_rank}", rank=self.rank, category="comms" + ): + patch = mx.distributed.recv( + (batch_size, patch_len, hidden_dim), patch.dtype, self.prev_rank, group=self.group, ) - mx.eval(encoder_hidden_states) + mx.eval(patch) + + if patch_idx == 0: + with trace( + name=f"recv {self.prev_rank}", rank=self.rank, category="comms" + ): + encoder_hidden_states = mx.distributed.recv( + (batch_size, text_seq_len, hidden_dim), + patch.dtype, + self.prev_rank, + group=self.group, + ) + mx.eval(encoder_hidden_states) if self.is_first_stage: patch, encoder_hidden_states = self.adapter.compute_embeddings( @@ -909,14 +964,22 @@ class DiffusionRunner: assert self.joint_block_wrappers is not None assert encoder_hidden_states is not None - for wrapper in self.joint_block_wrappers: - wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token) - encoder_hidden_states, patch = wrapper( - hidden_states=patch, - encoder_hidden_states=encoder_hidden_states, - text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_embeddings, - ) + with trace( + name=f"joint patch {patch_idx}", + rank=self.rank, + category="compute", + ): + for wrapper in self.joint_block_wrappers: + wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token) + encoder_hidden_states, patch = wrapper( + hidden_states=patch, + encoder_hidden_states=encoder_hidden_states, + text_embeddings=text_embeddings, + rotary_embeddings=image_rotary_embeddings, + ) + + if EXO_TRACING_ENABLED: + mx.eval(encoder_hidden_states, patch) if self.owns_concat_stage: assert encoder_hidden_states is not None @@ -925,49 +988,70 @@ class DiffusionRunner: if self.has_single_blocks or self.is_last_stage: patch = patch_concat else: - patch_concat = mx.distributed.send( - patch_concat, self.next_rank, group=self.group - ) - mx.async_eval(patch_concat) + with trace( + name=f"send {self.next_rank}", rank=self.rank, category="comms" + ): + patch_concat = mx.distributed.send( + patch_concat, self.next_rank, group=self.group + ) + mx.async_eval(patch_concat) elif self.has_joint_blocks and not self.is_last_stage: - patch = mx.distributed.send(patch, self.next_rank, group=self.group) - mx.async_eval(patch) + with trace(name=f"send {self.next_rank}", rank=self.rank, category="comms"): + patch = mx.distributed.send(patch, self.next_rank, group=self.group) + mx.async_eval(patch) if patch_idx == 0: assert encoder_hidden_states is not None - encoder_hidden_states = mx.distributed.send( - encoder_hidden_states, self.next_rank, group=self.group - ) - mx.async_eval(encoder_hidden_states) + with trace( + name=f"send {self.next_rank}", rank=self.rank, category="comms" + ): + encoder_hidden_states = mx.distributed.send( + encoder_hidden_states, self.next_rank, group=self.group + ) + mx.async_eval(encoder_hidden_states) if self.has_single_blocks: if not self.owns_concat_stage and not self.is_first_stage: patch_len = patch.shape[1] - patch = mx.distributed.recv( - (batch_size, text_seq_len + patch_len, hidden_dim), - patch.dtype, - self.prev_rank, - group=self.group, - ) - mx.eval(patch) + with trace( + name=f"recv {self.prev_rank}", rank=self.rank, category="comms" + ): + patch = mx.distributed.recv( + (batch_size, text_seq_len + patch_len, hidden_dim), + patch.dtype, + self.prev_rank, + group=self.group, + ) + mx.eval(patch) assert self.single_block_wrappers is not None - for wrapper in self.single_block_wrappers: - wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token) - patch = wrapper( - hidden_states=patch, - text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_embeddings, - ) + with trace( + name=f"single patch {patch_idx}", + rank=self.rank, + category="compute", + ): + for wrapper in self.single_block_wrappers: + wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token) + patch = wrapper( + hidden_states=patch, + text_embeddings=text_embeddings, + rotary_embeddings=image_rotary_embeddings, + ) + + if EXO_TRACING_ENABLED: + mx.eval(patch) if not self.is_last_stage: - patch = mx.distributed.send(patch, self.next_rank, group=self.group) - mx.async_eval(patch) + with trace( + name=f"send {self.next_rank}", rank=self.rank, category="comms" + ): + patch = mx.distributed.send(patch, self.next_rank, group=self.group) + mx.async_eval(patch) noise: mx.array | None = None if self.is_last_stage: - patch_img_only = patch[:, text_seq_len:, :] - noise = self.adapter.final_projection(patch_img_only, text_embeddings) + patch = patch[:, text_seq_len:, :] + noise = self.adapter.final_projection(patch, text_embeddings) return noise, encoder_hidden_states diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 0bc2e846..0d012fb6 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -16,8 +16,9 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs] ) from pydantic import ValidationError -from exo.shared.constants import EXO_MAX_CHUNK_SIZE +from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED from exo.shared.models.model_cards import ModelId, ModelTask +from exo.shared.tracing import clear_trace_buffer, get_trace_buffer from exo.shared.types.api import ImageGenerationStats from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk from exo.shared.types.common import CommandId @@ -27,6 +28,8 @@ from exo.shared.types.events import ( RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, + TraceEventData, + TracesCollected, ) from exo.shared.types.tasks import ( ConnectToGroup, @@ -408,6 +411,10 @@ def main( ) ) raise + finally: + _send_traces_if_enabled( + event_sender, task.task_id, shard_metadata.device_rank + ) current_status = RunnerReady() logger.info("runner ready") @@ -466,6 +473,10 @@ def main( ) ) raise + finally: + _send_traces_if_enabled( + event_sender, task.task_id, shard_metadata.device_rank + ) current_status = RunnerReady() logger.info("runner ready") @@ -640,6 +651,36 @@ def _send_image_chunk( ) +def _send_traces_if_enabled( + event_sender: MpSender[Event], + task_id: TaskId, + rank: int, +) -> None: + if not EXO_TRACING_ENABLED: + return + + traces = get_trace_buffer() + if traces: + trace_data = [ + TraceEventData( + name=t.name, + start_us=t.start_us, + duration_us=t.duration_us, + rank=t.rank, + category=t.category, + ) + for t in traces + ] + event_sender.send( + TracesCollected( + task_id=task_id, + rank=rank, + traces=trace_data, + ) + ) + clear_trace_buffer() + + def _process_image_response( response: ImageGenerationResponse | PartialImageResponse, command_id: CommandId,