← back to Exo
full mlx caching implementation
1b8b456ced1937b2c9140da8c07772b457304a7a · 2025-08-26 09:15:08 -0700 · Matt Beton
Files touched
M pyproject.tomlM run.shM scripts/src/exo_scripts/read_events.pyM src/exo/engines/mlx/utils_mlx.pyM src/exo/shared/logging.pyM src/exo/worker/runner/communication.pyM src/exo/worker/runner/runner.pyM src/exo/worker/runner/runner_supervisor.pyM src/exo/worker/runner/utils.pyM src/exo/worker/tests/test_integration/integration_utils.pyM src/exo/worker/tests/test_multimodel/test_inference_llama70B.pyM uv.lock
Diff
commit 1b8b456ced1937b2c9140da8c07772b457304a7a
Author: Matt Beton <matthew.beton@gmail.com>
Date: Tue Aug 26 09:15:08 2025 -0700
full mlx caching implementation
---
pyproject.toml | 1 +
run.sh | 12 +-
scripts/src/exo_scripts/read_events.py | 903 ++++++++++-----------
src/exo/engines/mlx/utils_mlx.py | 11 +-
src/exo/shared/logging.py | 4 +-
src/exo/worker/runner/communication.py | 11 +-
src/exo/worker/runner/runner.py | 2 +-
src/exo/worker/runner/runner_supervisor.py | 37 +-
src/exo/worker/runner/utils.py | 1 +
.../tests/test_integration/integration_utils.py | 1 +
.../test_multimodel/test_inference_llama70B.py | 80 +-
uv.lock | 68 ++
12 files changed, 581 insertions(+), 550 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 788405ff..52e708e2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
"transformers>=4.55.2",
"cobs>=1.2.2",
"loguru>=0.7.3",
+ "textual>=5.3.0",
]
[project.scripts]
diff --git a/run.sh b/run.sh
index 82e29cf1..8f329855 100755
--- a/run.sh
+++ b/run.sh
@@ -33,16 +33,16 @@ fi
# Configure MLX
# ./configure_mlx.sh
-# First command (worker) - changes based on replica flag
+# Second command (master) - changes based on replica flag
if [ "$REPLICA" = true ]; then
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export EXO_HOME=.exo; uv run exo-worker'\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true EXO_RUN_AS_REPLICA=1 EXO_HOME=.exo API_PORT=8001; uv run exo-master'\""
else
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c uv run exo-worker\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true; uv run exo-master'\""
fi
-# Second command (master) - changes based on replica flag
+# First command (worker) - changes based on replica flag
if [ "$REPLICA" = true ]; then
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true EXO_RUN_AS_REPLICA=1 EXO_HOME=.exo API_PORT=8001; uv run exo-master'\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export EXO_HOME=.exo; uv run exo-worker'\""
else
- osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true; uv run exo-master'\""
+ osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c uv run exo-worker\""
fi
\ No newline at end of file
diff --git a/scripts/src/exo_scripts/read_events.py b/scripts/src/exo_scripts/read_events.py
index 68fc9398..59493b26 100644
--- a/scripts/src/exo_scripts/read_events.py
+++ b/scripts/src/exo_scripts/read_events.py
@@ -1,26 +1,36 @@
-# pyright: reportAny=false
-
import asyncio
-import curses
-import time
import json
import argparse
import sys
-from logging import Logger
+import time
+from dataclasses import is_dataclass, asdict
+from logging import getLogger
from typing import List, Optional, Any, Sequence, Tuple
+# Your existing imports — unchanged
from exo.shared.types.state import State
from exo.shared.apply import apply
from exo.shared.db.sqlite.event_log_manager import EventLogManager, EventLogConfig
from exo.shared.types.events.components import EventFromEventLog
from exo.shared.types.events import Event
-# Globals
-logger: Logger = Logger('helper_log')
-event_log_manager: Optional[EventLogManager] = None
-worker_mode: bool = False
+# --- Third-party UI (new) ---
+from rich.syntax import Syntax
+from rich.text import Text
+from rich.panel import Panel
+from rich.console import RenderableType
-# Worker-related event types
+from textual.app import App, ComposeResult
+from textual.containers import Horizontal, Vertical
+from textual.widgets import Static, ListView, ListItem, Input, Footer, Label
+from textual.reactive import reactive
+from textual import on
+from textual.binding import Binding
+from textual.message import Message
+
+logger = getLogger("helper_log")
+
+# Worker-related event types (same set)
WORKER_EVENT_TYPES = {
'TaskCreated', 'TaskStateUpdated', 'TaskFailed', 'TaskDeleted',
'ChunkGenerated',
@@ -29,17 +39,19 @@ WORKER_EVENT_TYPES = {
}
+# ---------- Data / DB helpers (mostly your original logic) ----------
+
+event_log_manager: Optional[EventLogManager] = None
+
async def init_db() -> None:
global event_log_manager
- event_log_manager = EventLogManager(EventLogConfig(), logger)
+ event_log_manager = EventLogManager(EventLogConfig())
await event_log_manager.initialize()
-
async def get_events_since(since: int) -> Sequence[EventFromEventLog[Event]]:
- assert event_log_manager is not None
+ # type: ignore[attr-defined, return-value]
return await event_log_manager.global_events.get_events_since(since)
-
async def load_all_events() -> List[EventFromEventLog[Event]]:
events: List[EventFromEventLog[Event]] = []
since = 0
@@ -51,7 +63,6 @@ async def load_all_events() -> List[EventFromEventLog[Event]]:
since += len(new_events)
return events
-
def compute_states(events: List[EventFromEventLog[Event]]) -> List[State]:
states: List[State] = [State()]
state = states[0]
@@ -60,34 +71,95 @@ def compute_states(events: List[EventFromEventLog[Event]]) -> List[State]:
states.append(state)
return states
+def filter_worker_state(state: State) -> dict:
+ state_dict = json.loads(state.model_dump_json())
+ return {
+ 'node_status': state_dict.get('node_status', {}),
+ 'instances': state_dict.get('instances', {}),
+ 'runners': state_dict.get('runners', {}),
+ 'tasks': state_dict.get('tasks', {}),
+ 'last_event_applied_idx': state_dict.get('last_event_applied_idx', 0)
+ }
+
+def event_type_name(e: EventFromEventLog[Event]) -> str:
+ return type(e.event).__name__
+
+def is_worker_event(e: EventFromEventLog[Event]) -> bool:
+ return event_type_name(e) in WORKER_EVENT_TYPES
+
+def safe_json(obj: Any) -> str:
+ """Serialize unknown objects to JSON-ish string safely."""
+ def to_serializable(x: Any):
+ try:
+ if is_dataclass(x):
+ return asdict(x)
+ except Exception:
+ pass
+ if isinstance(x, (str, int, float, bool)) or x is None:
+ return x
+ if isinstance(x, dict):
+ return {str(k): to_serializable(v) for k, v in x.items()}
+ if isinstance(x, (list, tuple, set)):
+ return [to_serializable(v) for v in x]
+ try:
+ json.dumps(x) # type: ignore
+ return x
+ except Exception:
+ return repr(x)
+ try:
+ return json.dumps(to_serializable(obj), indent=2, ensure_ascii=False)
+ except Exception:
+ # Last resort
+ return repr(obj)
+
+def summarize_event_line(e: EventFromEventLog[Event], max_len: int = 160) -> Text:
+ etype = event_type_name(e)
+ attrs = vars(e.event)
+ prefix = Text(f"[{e.idx_in_log}] ", style="bold dim")
+ t = Text(etype, style="bold cyan")
+ t = prefix + t + Text(": ", style="dim")
+ first = True
+ for k, v in attrs.items():
+ if not first:
+ t.append(", ", style="dim")
+ first = False
+ t.append(str(k), style="magenta")
+ t.append("=")
+ # Coarse coloring by type
+ if isinstance(v, str):
+ t.append(repr(v), style="green")
+ elif isinstance(v, (int, float)):
+ t.append(repr(v), style="yellow")
+ elif isinstance(v, bool):
+ t.append(repr(v), style="cyan")
+ else:
+ t.append(repr(v), style="")
+ if len(t.plain) > max_len:
+ t.truncate(max_len - 1)
+ t.append("…", style="dim")
+ return t
+
+def event_detail_renderable(e: EventFromEventLog[Event]) -> RenderableType:
+ payload = {
+ "idx_in_log": e.idx_in_log,
+ "event_type": event_type_name(e),
+ "attributes": vars(e.event)
+ }
+ return Syntax(safe_json(payload), "json", word_wrap=True)
-def print_event(event: EventFromEventLog[Event]) -> None:
- event_type_name = type(event.event).__name__
- event_type = event_type_name.replace('_', ' ').title()
- attributes = ', '.join(f"{key}={value!r}" for key,
- value in vars(event.event).items())
- print(f"[{event.idx_in_log}] {event_type}: {attributes}")
+# ---------- Non-TUI (stdout) mode, like your current script ----------
-async def non_tui_mode() -> None:
+async def run_non_tui(worker_mode: bool) -> None:
await init_db()
events = await load_all_events()
states = compute_states(events)
final_state = states[-1]
if worker_mode:
- filtered_events = [e for e in events if type(
- e.event).__name__ in WORKER_EVENT_TYPES]
+ filtered_events = [e for e in events if is_worker_event(e)]
events = filtered_events
- # Recompute states? But states are cumulative, so perhaps just print filtered events and full state, or filter state too.
- state_dict = json.loads(final_state.model_dump_json())
- filtered_state = {
- 'node_status': state_dict.get('node_status', {}),
- 'instances': state_dict.get('instances', {}),
- 'runners': state_dict.get('runners', {}),
- 'tasks': state_dict.get('tasks', {}),
- 'last_event_applied_idx': state_dict.get('last_event_applied_idx', 0)
- }
+ filtered_state = filter_worker_state(final_state)
print("Final State (filtered):")
print(json.dumps(filtered_state, indent=2))
else:
@@ -95,464 +167,345 @@ async def non_tui_mode() -> None:
print(final_state.model_dump_json(indent=2))
print("\nEvents:")
- for event in events:
- print_event(event)
+ for e in events:
+ etype = event_type_name(e)
+ attrs = ', '.join(f"{k}={value!r}" for k, value in vars(e.event).items())
+ print(f"[{e.idx_in_log}] {etype}: {attrs}")
-async def update_events(wrapped_events: List[EventFromEventLog[Event]], states: List[State],
- filtered_indices: Optional[List[int]] = None) -> bool:
- last_since = len(wrapped_events)
- new_wrapped = await get_events_since(last_since)
- if new_wrapped:
- last_len = len(wrapped_events)
- for nw in new_wrapped:
- state = states[-1]
- new_state = apply(state, nw)
- states.append(new_state)
- wrapped_events.extend(new_wrapped)
- if filtered_indices is not None:
- for k in range(last_len, len(wrapped_events)):
- if type(wrapped_events[k].event).__name__ in WORKER_EVENT_TYPES:
- filtered_indices.append(k)
- return True
- return False
+# ---------- Textual TUI ----------
-
-def draw_state(win: Any, state: State, height: int, width: int, worker_mode: bool, state_scroll: int) -> int:
- win.clear()
- state_dict = json.loads(state.model_dump_json())
- if worker_mode:
- filtered_state = {
- 'node_status': state_dict.get('node_status', {}),
- 'instances': state_dict.get('instances', {}),
- 'runners': state_dict.get('runners', {}),
- 'tasks': state_dict.get('tasks', {}),
- 'last_event_applied_idx': state_dict.get('last_event_applied_idx', 0)
- }
- state_pretty = json.dumps(filtered_state, indent=2)
- else:
- state_pretty = json.dumps(state_dict, indent=2)
- lines = state_pretty.split('\n')
- max_scroll = max(0, len(lines) - height)
- current_scroll = min(state_scroll, max_scroll)
- for i in range(height):
- line_idx = current_scroll + i
- if line_idx >= len(lines):
- break
- line = lines[line_idx]
- y = i
- x = 0
- leading_spaces = len(line) - len(line.lstrip())
- win.addstr(y, x, ' ' * leading_spaces)
- x += leading_spaces
- stripped = line.lstrip()
- if stripped.startswith('"'):
- end_key = stripped.find('": ')
- if end_key != -1:
- key_str = stripped[:end_key + 3] # include ":
- win.addstr(y, x, key_str, curses.color_pair(3))
- x += len(key_str)
- value_str = stripped[end_key + 3:]
- if value_str.startswith('"'):
- color = 2
- elif value_str.replace('.', '', 1).isdigit() or (
- value_str.startswith('-') and value_str[1:].replace('.', '', 1).isdigit()):
- color = 4
- elif value_str in ['true', 'false', 'null']:
- color = 5
- elif value_str.startswith('{') or value_str.startswith('[') or value_str.startswith(
- '}') or value_str.startswith(']'):
- color = 0
- else:
- color = 0
- win.addstr(y, x, value_str, curses.color_pair(color))
- else:
- win.addstr(y, x, stripped)
- else:
- win.addstr(y, x, stripped)
- win.refresh()
- return current_scroll
-
-
-def get_event_pairs(event: EventFromEventLog[Event]) -> List[Tuple[str, int]]:
- pairs: List[Tuple[str, int]] = []
- idx_str = f"[{event.idx_in_log}] "
- pairs.append((idx_str, 5))
- event_type_name = type(event.event).__name__
- event_type = event_type_name.replace('_', ' ').title()
- pairs.append((event_type, 1))
- pairs.append((": ", 0))
- attrs = vars(event.event)
- first = True
- for key, value in attrs.items():
- if not first:
- pairs.append((", ", 0))
- first = False
- pairs.append((key, 3))
- pairs.append(("=", 0))
- v_str = repr(value)
- if isinstance(value, str):
- color = 2
- elif isinstance(value, (int, float)):
- color = 4
- elif isinstance(value, bool):
- color = 5
+class StateView(Static):
+ """Left pane: shows state JSON, with optional worker filter."""
+ def update_state(self, state: State, worker_mode: bool, index_in_log_for_status: Optional[int]) -> None:
+ if worker_mode:
+ data = filter_worker_state(state)
+ json_str = json.dumps(data, indent=2, ensure_ascii=False)
else:
- color = 6
- pairs.append((v_str, color))
- return pairs
-
-
-def calculate_event_lines(pairs: List[Tuple[str, int]], win_width: int, subsequent_indent: int) -> int:
- lines = 1
- x = 0
- for text, _ in pairs:
- i = 0
- while i < len(text):
- remaining = win_width - x
- part_len = min(len(text) - i, remaining)
- i += part_len
- x += part_len
- if i < len(text):
- lines += 1
- x = subsequent_indent
- return lines
-
-
-def render_event(win: Any, start_y: int, pairs: List[Tuple[str, int]], is_bold: bool, win_width: int,
- subsequent_indent: int) -> int:
- y = start_y
- x = 0
- for text, color in pairs:
- attr = curses.color_pair(color) | (curses.A_BOLD if is_bold else 0)
- i = 0
- while i < len(text):
- remaining = win_width - x
- part_len = min(len(text) - i, remaining)
- part = text[i:i + part_len]
- try:
- win.addstr(y, x, part, attr)
- except curses.error:
- pass
- i += part_len
- x += part_len
- if i < len(text):
- y += 1
- if y >= win.getmaxyx()[0]:
- return y
- x = subsequent_indent
- if x > 0:
- y += 1
- return y
-
-
-def draw_events(win: Any, events_list: List[EventFromEventLog[Event]], current_events: int, height: int) -> None:
- win.clear()
- if len(events_list) == 0:
- win.addstr(0, 0, "No events")
- win.refresh()
- return
- win_width = win.getmaxyx()[1]
- current_event = events_list[current_events]
- current_pairs = get_event_pairs(current_event)
- subsequent_indent = len(f"[{current_event.idx_in_log}] ")
- lines_current = calculate_event_lines(
- current_pairs, win_width, subsequent_indent)
- if lines_current > height:
- render_event(win, 0, current_pairs, True, win_width, subsequent_indent)
- win.refresh()
- return
-
- target_above = (height - lines_current) // 2
- target_below = height - lines_current - target_above
-
- # Collect previous events
- prev_events: List[int] = []
- remaining = target_above
- i = current_events - 1
- while i >= 0 and remaining > 0:
- event = events_list[i]
- pairs = get_event_pairs(event)
- indent = len(f"[{event.idx_in_log}] ")
- lines = calculate_event_lines(pairs, win_width, indent)
- if lines <= remaining:
- remaining -= lines
- prev_events.append(i)
- i -= 1
+ json_str = state.model_dump_json(indent=2)
+ syntax = Syntax(json_str, "json", word_wrap=True)
+ title = f"State after event #{index_in_log_for_status}" if index_in_log_for_status is not None else "Initial State"
+ self.update(Panel(syntax, title=title, border_style="cyan"))
+
+class EventListItem(ListItem):
+ def __init__(self, e: EventFromEventLog[Event]) -> None:
+ super().__init__(Static(summarize_event_line(e)))
+ self._event = e
+
+ @property
+ def wrapped_event(self) -> EventFromEventLog[Event]:
+ return self._event
+
+class EventDetail(Static):
+ """Right-bottom: details of the selected event."""
+ def show_event(self, e: Optional[EventFromEventLog[Event]]) -> None:
+ if e is None:
+ self.update(Panel(Text("No event selected.", style="dim"), title="Event Details"))
else:
- break
- prev_events.reverse()
-
- # Collect next events
- next_events: List[int] = []
- remaining = target_below
- j = current_events + 1
- while j < len(events_list) and remaining > 0:
- event = events_list[j]
- pairs = get_event_pairs(event)
- indent = len(f"[{event.idx_in_log}] ")
- lines = calculate_event_lines(pairs, win_width, indent)
- if lines <= remaining:
- remaining -= lines
- next_events.append(j)
- j += 1
+ self.update(Panel(event_detail_renderable(e), title=f"Event #{e.idx_in_log} • {event_type_name(e)}", border_style="magenta"))
+
+class StatusBar(Static):
+ def set_status(self, realtime: bool, total_events: int, current_idx_in_log: Optional[int]) -> None:
+ mode = "Realtime" if realtime else "Timetravel"
+ parts = [
+ f"[{mode}]",
+ f"Events: {total_events}",
+ ]
+ if current_idx_in_log is not None:
+ parts.append(f"Current: #{current_idx_in_log}")
+ parts.append("Keys: ↑/↓ Select • PgUp/PgDn Scroll • Ctrl+↑/↓ ±5 • [/] State PgUp/PgDn • g Goto • r Realtime • q Quit")
+ self.update(Text(" ".join(parts), style="dim"))
+
+
+class GotoPrompt(Static):
+ """Simple inline goto prompt (appears above Footer)."""
+ class Submitted(Message):
+ def __init__(self, value: Optional[int]) -> None:
+ super().__init__()
+ self.value = value
+
+ def compose(self) -> ComposeResult:
+ yield Label("Go to event id (idx_in_log):", id="goto-label")
+ yield Input(placeholder="e.g., 123", id="goto-input")
+
+ def on_mount(self) -> None:
+ self.query_one(Input).focus()
+
+ @on(Input.Submitted)
+ def _submitted(self, event: Input.Submitted) -> None:
+ text = (event.value or "").strip()
+ try:
+ value = int(text)
+ except ValueError:
+ value = None
+ self.post_message(self.Submitted(value))
+
+
+class EventLogApp(App):
+ CSS = """
+ Screen {
+ layout: vertical;
+ }
+ #main {
+ height: 1fr;
+ }
+ #left {
+ width: 60%;
+ }
+ #right {
+ width: 40%;
+ }
+ #events {
+ height: 3fr;
+ }
+ #detail {
+ height: 2fr;
+ border: tall;
+ }
+ #status {
+ height: 1;
+ padding: 0 1;
+ }
+ #goto {
+ dock: bottom;
+ height: 3;
+ padding: 1 2;
+ background: $panel;
+ border: round $accent;
+ }
+ """
+
+ BINDINGS = [
+ Binding("q", "quit", "Quit"),
+ Binding("r", "toggle_realtime", "Realtime"),
+ Binding("[", "state_page_up", "State PgUp"),
+ Binding("]", "state_page_down", "State PgDn"),
+ Binding("g", "prompt_goto", "Goto"),
+ Binding("ctrl+up", "jump_up", "Jump Up"),
+ Binding("ctrl+down", "jump_down", "Jump Down"),
+ ]
+
+ # Reactive state
+ realtime: reactive[bool] = reactive(False)
+ worker_mode: bool
+
+ # Data
+ wrapped_events: List[EventFromEventLog[Event]]
+ states: List[State]
+ filtered_indices: Optional[List[int]] # maps filtered idx -> original idx
+ update_interval: float = 1.0
+ _poll_timer = None
+
+ def __init__(self, worker_mode: bool) -> None:
+ super().__init__()
+ self.worker_mode = worker_mode
+ self.wrapped_events = []
+ self.states = [State()]
+ self.filtered_indices = None
+
+ async def on_mount(self) -> None:
+ await init_db()
+ await self._initial_load()
+ # periodic polling for new events
+ self._poll_timer = self.set_interval(self.update_interval, self._tick_poll)
+ # Put list selection at end (last event) by default
+ self._select_last()
+
+ async def _initial_load(self) -> None:
+ self.wrapped_events = await load_all_events()
+ self.states = compute_states(self.wrapped_events)
+
+ # Build filtered view if needed
+ if self.worker_mode:
+ self.filtered_indices = [i for i, e in enumerate(self.wrapped_events) if is_worker_event(e)]
else:
- break
-
- # Calculate total lines
- total_lines = lines_current
- for idx in prev_events:
- event = events_list[idx]
- pairs = get_event_pairs(event)
- indent = len(f"[{event.idx_in_log}] ")
- total_lines += calculate_event_lines(pairs, win_width, indent)
- for idx in next_events:
- event = events_list[idx]
- pairs = get_event_pairs(event)
- indent = len(f"[{event.idx_in_log}] ")
- total_lines += calculate_event_lines(pairs, win_width, indent)
-
- padding = (height - total_lines) // 2 if total_lines < height else 0
-
- y = padding
- # Draw prev
- for idx in prev_events:
- event = events_list[idx]
- pairs = get_event_pairs(event)
- indent = len(f"[{event.idx_in_log}] ")
- y = render_event(win, y, pairs, False, win_width, indent)
-
- # Draw current
- y = render_event(win, y, current_pairs, True, win_width, subsequent_indent)
-
- # Draw next
- for idx in next_events:
- event = events_list[idx]
- pairs = get_event_pairs(event)
- indent = len(f"[{event.idx_in_log}] ")
- y = render_event(win, y, pairs, False, win_width, indent)
-
- win.refresh()
-
-
-def draw_status(win: Any, realtime: bool, current: int, total_events: int) -> None:
- win.clear()
- mode = "Realtime" if realtime else "Timetravel"
- win.addstr(0, 0,
- f"Mode: {mode} | Current event: {current} / {total_events} | Arrows: navigate events, [/]: scroll state, g: goto, r: toggle realtime, q: quit")
- win.refresh()
-
-
-def get_input(stdscr: Any, prompt: str) -> str:
- curses.echo()
- stdscr.addstr(0, 0, prompt)
- stdscr.refresh()
- input_str = stdscr.getstr(0, len(prompt), 20).decode('utf-8')
- curses.noecho()
- return input_str
-
-
-def get_key(win: Any) -> Any:
- ch = win.getch()
- if ch == -1:
- return -1
- if ch == 27:
- ch2 = win.getch()
- if ch2 == -1:
- return 27
- if ch2 == 91:
- ch3 = win.getch()
- if ch3 == -1:
+ self.filtered_indices = None
+
+ # Populate the ListView
+ lv = self.query_one("#events", ListView)
+ lv.clear()
+ events_to_show = self._view_events()
+ for e in events_to_show:
+ lv.append(EventListItem(e))
+
+ # Update left state & details
+ self._refresh_views()
+
+ def compose(self) -> ComposeResult:
+ # Layout: [Header optional] -> main Horizontal -> Status bar + Footer
+ with Horizontal(id="main"):
+ with Vertical(id="left"):
+ yield StateView(id="state")
+ with Vertical(id="right"):
+ yield ListView(id="events")
+ yield EventDetail(id="detail")
+ yield StatusBar(id="status")
+ yield Footer()
+
+ def _current_original_index(self) -> int:
+ lv = self.query_one("#events", ListView)
+ idx = lv.index
+ if idx is None or idx < 0:
+ return -1
+ if self.filtered_indices is not None:
+ if idx >= len(self.filtered_indices):
return -1
- if ch3 == 65:
- return curses.KEY_UP
- if ch3 == 66:
- return curses.KEY_DOWN
- if ch3 == 53:
- ch4 = win.getch()
- if ch4 == 126:
- return curses.KEY_PPAGE
- if ch3 == 54:
- ch4 = win.getch()
- if ch4 == 126:
- return curses.KEY_NPAGE
- if ch3 == 49:
- ch4 = win.getch()
- if ch4 == -1:
- return -1
- if ch4 == 59:
- ch5 = win.getch()
- if ch5 == -1:
- return -1
- if ch5 == 53:
- ch6 = win.getch()
- if ch6 == -1:
- return -1
- if ch6 == 65:
- return 'CTRL_UP'
- if ch6 == 66:
- return 'CTRL_DOWN'
- return ch
-
-
-def tui(stdscr: Any) -> None:
- curses.start_color()
- curses.init_pair(1, curses.COLOR_BLUE, curses.COLOR_BLACK)
- curses.init_pair(2, curses.COLOR_GREEN, curses.COLOR_BLACK)
- curses.init_pair(3, curses.COLOR_MAGENTA, curses.COLOR_BLACK)
- curses.init_pair(4, curses.COLOR_YELLOW, curses.COLOR_BLACK)
- curses.init_pair(5, curses.COLOR_CYAN, curses.COLOR_BLACK)
- curses.init_pair(6, curses.COLOR_WHITE, curses.COLOR_BLACK)
- curses.use_default_colors()
- stdscr.timeout(100)
- curses.curs_set(0)
-
- wrapped_events: List[EventFromEventLog[Event]] = []
- states: List[State] = [State()]
- asyncio.run(init_db())
- asyncio.run(update_events(wrapped_events, states)) # Initial load
+ return self.filtered_indices[idx]
+ return idx
+
+ def _view_events(self) -> List[EventFromEventLog[Event]]:
+ if self.filtered_indices is not None:
+ return [self.wrapped_events[i] for i in self.filtered_indices]
+ return self.wrapped_events
+
+ def _select_last(self) -> None:
+ lv = self.query_one("#events", ListView)
+ n = len(lv.children)
+ if n:
+ lv.index = n - 1
+
+ def _refresh_views(self) -> None:
+ # Update State pane and Detail pane and Status bar
+ original_idx = self._current_original_index()
+ state_idx = (original_idx + 1) if original_idx >= 0 else 0
+ state = self.states[state_idx]
+ state_view = self.query_one("#state", StateView)
+ idx_in_log = None
+ if original_idx >= 0:
+ idx_in_log = self.wrapped_events[original_idx].idx_in_log
+ state_view.update_state(state, self.worker_mode, idx_in_log)
+
+ # Detail pane
+ detail = self.query_one("#detail", EventDetail)
+ current_event = self.wrapped_events[original_idx] if original_idx >= 0 else None
+ detail.show_event(current_event)
+
+ # Status bar
+ status = self.query_one("#status", StatusBar)
+ total_events = len(self.wrapped_events)
+ status.set_status(self.realtime, total_events, current_event.idx_in_log if current_event else None)
+
+ async def _poll_once(self) -> bool:
+ """Fetch and append new events; return True if updated."""
+ last_since = len(self.wrapped_events)
+ new_wrapped = await get_events_since(last_since)
+ if not new_wrapped:
+ return False
+
+ # Extend states incrementally (avoid recomputing all)
+ for nw in new_wrapped:
+ state = self.states[-1]
+ self.states.append(apply(state, nw))
+
+ start_len = len(self.wrapped_events)
+ self.wrapped_events.extend(new_wrapped)
+
+ # Update filtered mapping and UI list
+ lv = self.query_one("#events", ListView)
+ if self.worker_mode:
+ if self.filtered_indices is None:
+ self.filtered_indices = []
+ for k in range(start_len, len(self.wrapped_events)):
+ if is_worker_event(self.wrapped_events[k]):
+ self.filtered_indices.append(k)
+ lv.append(EventListItem(self.wrapped_events[k]))
+ else:
+ for k in range(start_len, len(self.wrapped_events)):
+ lv.append(EventListItem(self.wrapped_events[k]))
- filtered_indices: Optional[List[int]] = None
- current_filtered: int = -1
- current: int = -1
- if worker_mode:
- filtered_indices = [i for i in range(len(wrapped_events)) if
- type(wrapped_events[i].event).__name__ in WORKER_EVENT_TYPES]
- current_filtered = len(filtered_indices) - \
- 1 if filtered_indices else -1
- else:
- current = len(wrapped_events) - 1 if wrapped_events else -1
+ # Auto-follow the tail in realtime mode
+ if self.realtime:
+ self._select_last()
- realtime: bool = False
- last_update: float = time.time()
- update_interval: float = 1.0
- state_scroll: int = 0
+ # Refresh panes
+ self._refresh_views()
+ return True
- while True:
- height, width = stdscr.getmaxyx()
- status_height = 1
- pane_height = height - status_height
- pane_width = width // 2
+ def _tick_poll(self) -> None:
+ # called by timer; schedule the async poll
+ asyncio.create_task(self._poll_once())
+
+ # ------ Actions / key handlers ------
+ def action_quit(self) -> None:
+ self.exit()
+
+ def action_toggle_realtime(self) -> None:
+ self.realtime = not self.realtime
+ if self.realtime:
+ self._select_last()
+ self._refresh_views()
+
+ def action_state_page_up(self) -> None:
+ state_view = self.query_one("#state", StateView)
+ state_view.scroll_page_up()
+
+ def action_state_page_down(self) -> None:
+ state_view = self.query_one("#state", StateView)
+ state_view.scroll_page_down()
+
+ def action_jump_up(self) -> None:
+ lv = self.query_one("#events", ListView)
+ if lv.children:
+ lv.index = max(0, (lv.index or 0) - 5)
+ self._refresh_views()
+
+ def action_jump_down(self) -> None:
+ lv = self.query_one("#events", ListView)
+ if lv.children:
+ lv.index = min(len(lv.children) - 1, (lv.index or 0) + 5)
+ self._refresh_views()
+
+ def action_prompt_goto(self) -> None:
+ # mount a small prompt near bottom
+ if self.query("#goto"):
+ return
+ prompt = GotoPrompt(id="goto")
+ self.mount(prompt)
+
+ @on(GotoPrompt.Submitted)
+ def _on_goto_submitted(self, msg: GotoPrompt.Submitted) -> None:
+ # Remove prompt
+ for node in self.query("#goto"):
+ node.remove()
+
+ if msg.value is None:
+ return
+
+ target = msg.value
+ # find in current view's idx_in_log
+ events_to_show = self._view_events()
+ lv = self.query_one("#events", ListView)
+ for i, e in enumerate(events_to_show):
+ if e.idx_in_log == target:
+ lv.index = i
+ self._refresh_views()
+ break
- state_win = curses.newwin(pane_height, pane_width, 0, 0)
- events_win = curses.newwin(
- pane_height, width - pane_width, 0, pane_width)
- status_win = curses.newwin(status_height, width, pane_height, 0)
+ @on(ListView.Highlighted, "#events")
+ @on(ListView.Selected, "#events")
+ def _on_event_selected(self, *_: Any) -> None:
+ # Update panes when selection changes
+ self._refresh_views()
- if worker_mode:
- assert filtered_indices is not None
- current_original = filtered_indices[current_filtered] if current_filtered >= 0 else -1
- events_list = [wrapped_events[i] for i in filtered_indices]
- current_events = current_filtered
- else:
- current_original = current
- events_list = wrapped_events
- current_events = current
-
- state_idx = current_original + 1 if current_original >= 0 else 0
- state_scroll = draw_state(
- state_win, states[state_idx], pane_height, pane_width, worker_mode, state_scroll)
- draw_events(events_win, events_list, current_events, pane_height)
- total_events = len(wrapped_events) - 1 if wrapped_events else -1
- draw_status(status_win, realtime,
- current_original if worker_mode else current, total_events)
-
- key = get_key(stdscr)
- if key != -1:
- if key == curses.KEY_UP:
- if worker_mode and current_filtered > 0:
- current_filtered -= 1
- elif not worker_mode and current > 0:
- current -= 1
- elif key == 'CTRL_UP':
- if worker_mode:
- current_filtered = max(0, current_filtered - 5)
- else:
- current = max(0, current - 5)
- elif key == curses.KEY_DOWN:
- assert filtered_indices is not None
- if worker_mode and current_filtered < len(filtered_indices) - 1:
- current_filtered += 1
- elif not worker_mode and current < len(wrapped_events) - 1:
- current += 1
- elif key == 'CTRL_DOWN':
- assert filtered_indices is not None
- if worker_mode:
- current_filtered = min(
- len(filtered_indices) - 1, current_filtered + 5)
- else:
- current = min(len(wrapped_events) - 1, current + 5)
- elif key == ord('['):
- state_scroll = max(0, state_scroll - pane_height // 2)
- elif key == ord(']'):
- state_scroll += pane_height // 2 # clamped in draw_state
- elif key == ord('q'):
- break
- elif key == ord('r'):
- realtime = not realtime
- if realtime:
- assert filtered_indices is not None
- if worker_mode:
- current_filtered = len(
- filtered_indices) - 1 if filtered_indices else -1
- else:
- current = len(wrapped_events) - \
- 1 if wrapped_events else -1
- state_scroll = 0
- elif key == ord('g'):
- stdscr.timeout(-1) # block for input
- input_str = get_input(status_win, "Go to event: ")
- try:
- goto = int(input_str)
- if worker_mode:
- assert filtered_indices is not None
- for i, orig in enumerate(filtered_indices):
- if wrapped_events[orig].idx_in_log == goto:
- current_filtered = i
- state_scroll = 0
- break
- else:
- for i in range(len(wrapped_events)):
- if wrapped_events[i].idx_in_log == goto:
- current = i
- state_scroll = 0
- break
- except ValueError:
- pass
- stdscr.timeout(100)
- status_win.clear()
- status_win.refresh()
-
- if realtime and time.time() - last_update > update_interval:
- updated = asyncio.run(update_events(
- wrapped_events, states, filtered_indices if worker_mode else None))
- if updated:
- assert filtered_indices is not None
- if worker_mode:
- current_filtered = len(filtered_indices) - 1
- else:
- current = len(wrapped_events) - 1
- state_scroll = 0
- last_update = time.time()
+# ---------- Entrypoint ----------
-if __name__ == "__main__":
- parser = argparse.ArgumentParser(
- description='Read and display events from the event log')
+def main() -> None:
+ parser = argparse.ArgumentParser(description='Read and display events from the event log (Textual UI)')
parser.add_argument('--worker', action='store_true',
help='Only show worker-related events (task, streaming, instance, runner status)')
+ parser.add_argument('--no-ui', action='store_true',
+ help='Print to stdout (non-interactive), like the original non-TUI mode')
args = parser.parse_args()
- worker_mode = args.worker
+ # Non-interactive fallback if no TTY or user requests it
+ if args.no_ui or not sys.stdout.isatty():
+ asyncio.run(run_non_tui(worker_mode=args.worker))
+ return
- if not sys.stdout.isatty():
- asyncio.run(non_tui_mode())
- else:
- try:
- curses.wrapper(tui)
- except curses.error as e:
- if "could not find terminal" in str(e):
- print("Error: Could not find terminal. Falling back to non-TUI mode.")
- asyncio.run(non_tui_mode())
- else:
- raise
+ # TUI mode
+ app = EventLogApp(worker_mode=args.worker)
+ app.run()
+
+if __name__ == "__main__":
+ main()
diff --git a/src/exo/engines/mlx/utils_mlx.py b/src/exo/engines/mlx/utils_mlx.py
index 6ac3dc6e..daf1636b 100644
--- a/src/exo/engines/mlx/utils_mlx.py
+++ b/src/exo/engines/mlx/utils_mlx.py
@@ -60,19 +60,16 @@ def mlx_setup(
target_cache = int(1.10 * (model_bytes + kv_bytes)) # +10% slack
target_cache = min(target_cache, int(cache_frac_of_mrwss * mrwss))
target_cache = min(target_cache, memsize)
- runner_print(f"{target_cache=}")
+ runner_print(f"{target_cache=}")
mx.set_cache_limit(max(target_cache, 0))
- return
-
- # Optional hard cap (keeps total MLX usage under control)
- with contextlib.suppress(Exception):
- mx.set_memory_limit(int(0.85 * mrwss))
# Wiring: off by default; if you re‑enable, wire at most a small fraction.
if wired_frac_of_mrwss > 0.0:
- target_wired = min(int(wired_frac_of_mrwss * mrwss), int(0.5 * model_bytes))
+ target_wired = int(wired_frac_of_mrwss * mrwss)
target_wired = min(target_wired, target_cache) # don’t wire more than cache
+
+ runner_print(f"{target_wired=}")
with contextlib.suppress(Exception): # older macOS won’t have this
mx.set_wired_limit(max(target_wired, 0))
diff --git a/src/exo/shared/logging.py b/src/exo/shared/logging.py
index 4946f1ad..2798ffbe 100644
--- a/src/exo/shared/logging.py
+++ b/src/exo/shared/logging.py
@@ -20,7 +20,7 @@ def logger_setup(log_file: Path, verbosity: int = 0):
if verbosity == 0:
_ = logger.add( # type: ignore
sys.__stderr__, # type: ignore
- format="[ {time:hh:mmA} | <level>{level: <8}</level>] <level>{message}</level>",
+ format="[ {time:hh:mm:ss.SSSSA} | <level>{level: <8}</level>] <level>{message}</level>",
level="INFO",
colorize=True,
enqueue=True,
@@ -29,7 +29,7 @@ def logger_setup(log_file: Path, verbosity: int = 0):
elif verbosity == 1:
_ = logger.add( # type: ignore
sys.__stderr__, # type: ignore
- format="[ {time:hh:mmA} | <level>{level: <8}</level>] <level>{message}</level>",
+ format="[ {time:hh:mm:ss.SSSSA} | <level>{level: <8}</level>] <level>{message}</level>",
level="INFO",
colorize=True,
enqueue=True,
diff --git a/src/exo/worker/runner/communication.py b/src/exo/worker/runner/communication.py
index 0b889aa4..d02ffb02 100644
--- a/src/exo/worker/runner/communication.py
+++ b/src/exo/worker/runner/communication.py
@@ -58,16 +58,17 @@ def runner_write_response(obj: RunnerResponse) -> None:
async def supervisor_read_response(
proc: asyncio.subprocess.Process,
-) -> RunnerResponse | None:
+) -> RunnerResponse:
assert proc.stdout is not None, (
"proc.stdout should not be None when created with stdout=PIPE"
)
- line_bytes: bytes = await asyncio.wait_for(proc.stdout.readline(), timeout=180)
+ # TODO: We could put a timeout on this if we decide to send heartbeats from the runner.
+ # This lets us handle cases where the process dies at some point not during an inference.
+ line_bytes: bytes = await proc.stdout.readline()
+ if not line_bytes:
+ raise EOFError('No more data to read when reading response from runner.')
line: str = line_bytes.decode("utf-8").strip()
- if not line:
- return None
-
try:
return RunnerResponseTypeAdapter.validate_json(line)
except Exception as err:
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 307de378..287f1e2a 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -112,7 +112,7 @@ async def main():
model_shard_meta = setup_message.model_shard_meta
hosts = setup_message.hosts
- mlx_setup(int(get_weights_size_kb(model_shard_meta) // 2**10))
+ mlx_setup(int(get_weights_size_kb(model_shard_meta) // 2**10), cache_frac_of_mrwss=0.8, wired_frac_of_mrwss=0.8)
# For testing - these are fake break conditions
if model_shard_meta.immediate_exception:
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index 87cfd7d7..665f00c4 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -114,6 +114,13 @@ class RunnerSupervisor:
"""
Read from the queue with a timeout, but also check if the read_task has failed.
"""
+ try:
+ assert not self.read_task.done()
+ except AssertionError as e_assert:
+ e = self.read_task.exception()
+ assert e is not None
+ raise e from e_assert
+
queue_task = asyncio.create_task(self.read_queue.get())
done, pending = await asyncio.wait(
@@ -137,13 +144,14 @@ class RunnerSupervisor:
return response
if self.read_task in done:
- await self.read_task # Re-raises any exception from read_task
- logger.error(
- "Unreachable code run. We should have raised an error on the read_task being done."
- )
-
+ try:
+ await self.read_task # Re-raises any exception from read_task
+ except Exception:
+ raise # bubble up exception
+ raise RunnerError("RunnerStopped", "Runner read loop terminated unexpectedly before any response.", "")
+
# if we haven't read from the queue, we have timed out.
- await self.astop()
+ await self.astop() # TODO: This could be handled by the called or _read_with_error_check - as we don't want a false Timeout to bring the whole runner down.
raise asyncio.TimeoutError()
async def stream_response(
@@ -186,7 +194,7 @@ class RunnerSupervisor:
try:
response = await self._read_with_error_check(timeout)
except asyncio.TimeoutError as e:
- logger.bind(user_facing=True).info(
+ logger.bind(user_facing=True).error(
f"Generation timed out during {'prefil' if timeout == prefil_timeout else 'decoding stage'}"
)
raise e
@@ -219,16 +227,17 @@ class RunnerSupervisor:
async def _read_coro(self):
while True:
- response: RunnerResponse | None = await supervisor_read_response(
- self.runner_process
- )
- if response is None:
- # Runner process died unexpectedly (C++ crash)
+ try:
+ response: RunnerResponse = await supervisor_read_response(
+ self.runner_process
+ )
+ except EOFError:
e = await self._raise_crashed()
if e:
- raise e from EOFError
+ # Runner process died unexpectedly (C++ crash)
+ raise e from EOFError # TODO: Do we just want to create an error and put it on the read_queue here?
else:
- break
+ continue
match response:
case PrintResponse():
diff --git a/src/exo/worker/runner/utils.py b/src/exo/worker/runner/utils.py
index 328d1a07..e3ddae62 100644
--- a/src/exo/worker/runner/utils.py
+++ b/src/exo/worker/runner/utils.py
@@ -68,6 +68,7 @@ def get_init_timeout(model_shard_meta: ShardMetadata) -> float:
def get_prefil_timeout(model_shard_meta: ShardMetadata) -> float:
+ return 30.0 # TODO: Proper prefil timeout calculation, but this requires knowing the number of tokens in the prompt.
weights_size_gb = get_weights_size_kb(model_shard_meta) / (1024 * 1024)
tokens = 1000 # constant for now - the prompt is only tokenized in the device...
diff --git a/src/exo/worker/tests/test_integration/integration_utils.py b/src/exo/worker/tests/test_integration/integration_utils.py
index 0654ad77..9d088a70 100644
--- a/src/exo/worker/tests/test_integration/integration_utils.py
+++ b/src/exo/worker/tests/test_integration/integration_utils.py
@@ -55,6 +55,7 @@ async def read_streaming_response(
event.chunk, TokenChunk
):
response_string += event.chunk.text
+ token_count += 1
if event.chunk.finish_reason:
finish_reason = event.chunk.finish_reason
diff --git a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
index 6e9ace7f..c71aafc8 100644
--- a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
+++ b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
@@ -183,15 +183,15 @@ async def test_ttft(
if not first_chunk_seen_1:
await asyncio.sleep(0.01)
- _, seen_task_finished_1, response_string_1, _ = await read_streaming_response(
+ _, seen_task_finished_1, response_string_1, token_count_1 = await read_streaming_response(
global_events
)
- # # total_time_1 = time.time() - task_created_time_1
+ total_time_1 = time.time() - task_created_time_1
assert seen_task_finished_1
# Wait for first task to complete
- await asyncio.sleep(3.0)
+ await asyncio.sleep(5.0)
# Second inference
task2_params = ChatCompletionTaskParams(
@@ -238,10 +238,10 @@ async def test_ttft(
if not first_chunk_seen_2:
await asyncio.sleep(0.01)
- _, seen_task_finished_2, response_string_2, _ = await read_streaming_response(
+ _, seen_task_finished_2, response_string_2, token_count_2 = await read_streaming_response(
global_events, filter_task=TASK_2_ID
)
- # # total_time_2 = time.time() - task_created_time_2
+ total_time_2 = time.time() - task_created_time_2
assert seen_task_finished_2
assert time_to_first_token_1
@@ -249,41 +249,41 @@ async def test_ttft(
# Calculate TPS metrics
# Prompt is approximately 45 tokens according to user
- # prompt_tokens = 45
-
- # # Prefill TPS = prompt tokens / time to first token
- # prefill_tps_1 = prompt_tokens / time_to_first_token_1 if time_to_first_token_1 > 0 else 0
- # prefill_tps_2 = prompt_tokens / time_to_first_token_2 if time_to_first_token_2 > 0 else 0
-
- # # Generation TPS = generated tokens / generation time
- # # Generation time = total time - time to first token
- # generation_time_1 = total_time_1 - time_to_first_token_1
- # generation_time_2 = total_time_2 - time_to_first_token_2
- # generation_tps_1 = token_count_1 / generation_time_1 if generation_time_1 > 0 else 0
- # generation_tps_2 = token_count_2 / generation_time_2 if generation_time_2 > 0 else 0
-
- # # Display time to first token profiling results
- # print("\n=== Time to First Token Profiling ===")
- # print(f"First inference ('{task1.task_params.messages[0].content}'):")
- # print(f" Time to first token: {time_to_first_token_1:.3f}s")
- # print(f" Total completion time: {total_time_1:.3f}s")
- # print(f" Tokens generated: {token_count_1}")
- # print(f" Response length: {len(response_string_1)} chars")
- # print(f" Prefill TPS: {prefill_tps_1:.1f} tokens/sec ({prompt_tokens} prompt tokens / {time_to_first_token_1:.3f}s)")
- # print(f" Generation TPS: {generation_tps_1:.1f} tokens/sec ({token_count_1} tokens / {generation_time_1:.3f}s)")
-
- # print(f"\nSecond inference ('{task2.task_params.messages[0].content}'):")
- # print(f" Time to first token: {time_to_first_token_2:.3f}s")
- # print(f" Total completion time: {total_time_2:.3f}s")
- # print(f" Tokens generated: {token_count_2}")
- # print(f" Response length: {len(response_string_2)} chars")
- # print(f" Prefill TPS: {prefill_tps_2:.1f} tokens/sec ({prompt_tokens} prompt tokens / {time_to_first_token_2:.3f}s)")
- # print(f" Generation TPS: {generation_tps_2:.1f} tokens/sec ({token_count_2} tokens / {generation_time_2:.3f}s)")
-
- # print("\nComparison:")
- # print(f" Second inference time to first token: {time_to_first_token_2/time_to_first_token_1:.2f}x the first")
- # print(f" Second inference prefill TPS: {prefill_tps_2/prefill_tps_1:.2f}x the first")
- # print(f" Second inference generation TPS: {generation_tps_2/generation_tps_1:.2f}x the first")
+ prompt_tokens = 45
+
+ # Prefill TPS = prompt tokens / time to first token
+ prefill_tps_1 = prompt_tokens / time_to_first_token_1 if time_to_first_token_1 > 0 else 0
+ prefill_tps_2 = prompt_tokens / time_to_first_token_2 if time_to_first_token_2 > 0 else 0
+
+ # Generation TPS = generated tokens / generation time
+ # Generation time = total time - time to first token
+ generation_time_1 = total_time_1 - time_to_first_token_1
+ generation_time_2 = total_time_2 - time_to_first_token_2
+ generation_tps_1 = token_count_1 / generation_time_1 if generation_time_1 > 0 else 0
+ generation_tps_2 = token_count_2 / generation_time_2 if generation_time_2 > 0 else 0
+
+ # Display time to first token profiling results
+ print("\n=== Time to First Token Profiling ===")
+ print(f"First inference ('{task1.task_params.messages[0].content}'):")
+ print(f" Time to first token: {time_to_first_token_1:.3f}s")
+ print(f" Total completion time: {total_time_1:.3f}s")
+ print(f" Tokens generated: {token_count_1}")
+ print(f" Response length: {len(response_string_1)} chars")
+ print(f" Prefill TPS: {prefill_tps_1:.1f} tokens/sec ({prompt_tokens} prompt tokens / {time_to_first_token_1:.3f}s)")
+ print(f" Generation TPS: {generation_tps_1:.1f} tokens/sec ({token_count_1} tokens / {generation_time_1:.3f}s)")
+
+ print(f"\nSecond inference ('{task2.task_params.messages[0].content}'):")
+ print(f" Time to first token: {time_to_first_token_2:.3f}s")
+ print(f" Total completion time: {total_time_2:.3f}s")
+ print(f" Tokens generated: {token_count_2}")
+ print(f" Response length: {len(response_string_2)} chars")
+ print(f" Prefill TPS: {prefill_tps_2:.1f} tokens/sec ({prompt_tokens} prompt tokens / {time_to_first_token_2:.3f}s)")
+ print(f" Generation TPS: {generation_tps_2:.1f} tokens/sec ({token_count_2} tokens / {generation_time_2:.3f}s)")
+
+ print("\nComparison:")
+ print(f" Second inference time to first token: {time_to_first_token_2/time_to_first_token_1:.2f}x the first")
+ print(f" Second inference prefill TPS: {prefill_tps_2/prefill_tps_1:.2f}x the first")
+ print(f" Second inference generation TPS: {generation_tps_2/generation_tps_1:.2f}x the first")
# Basic assertions to ensure responses make sense
assert len(response_string_1) > 0
diff --git a/uv.lock b/uv.lock
index 9abbbc8c..888d683e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -269,6 +269,7 @@ dependencies = [
{ name = "rustworkx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "sqlalchemy", extra = ["asyncio"], marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "sqlmodel", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "textual", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "typeguard", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "types-aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -313,6 +314,7 @@ requires-dist = [
{ name = "rustworkx", specifier = ">=0.17.1" },
{ name = "sqlalchemy", extras = ["asyncio"], specifier = ">=2.0.43" },
{ name = "sqlmodel", specifier = ">=0.0.24" },
+ { name = "textual", specifier = ">=5.3.0" },
{ name = "transformers", specifier = ">=4.55.2" },
{ name = "typeguard", specifier = ">=4.4.4" },
{ name = "types-aiofiles", specifier = ">=24.1.0.20250708" },
@@ -567,6 +569,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b3/4a/4175a563579e884192ba6e81725fc0448b042024419be8d83aa8a80a3f44/jiter-0.10.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3aa96f2abba33dc77f79b4cf791840230375f9534e5fac927ccceb58c5e604a5", size = 354213, upload-time = "2025-05-18T19:04:41.894Z" },
]
+[[package]]
+name = "linkify-it-py"
+version = "2.0.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "uc-micro-py", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/2a/ae/bb56c6828e4797ba5a4821eec7c43b8bf40f69cda4d4f5f8c8a2810ec96a/linkify-it-py-2.0.3.tar.gz", hash = "sha256:68cda27e162e9215c17d786649d1da0021a451bdc436ef9e0fa0ba5234b9b048", size = 27946, upload-time = "2024-02-04T14:48:04.179Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/04/1e/b832de447dee8b582cac175871d2f6c3d5077cc56d5575cadba1fd1cccfa/linkify_it_py-2.0.3-py3-none-any.whl", hash = "sha256:6bcbc417b0ac14323382aef5c5192c0075bf8a9d6b41820a2b66371eac6b6d79", size = 19820, upload-time = "2024-02-04T14:48:02.496Z" },
+]
+
[[package]]
name = "loguru"
version = "0.7.3"
@@ -588,6 +602,14 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" },
]
+[package.optional-dependencies]
+linkify = [
+ { name = "linkify-it-py", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+]
+plugins = [
+ { name = "mdit-py-plugins", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+]
+
[[package]]
name = "markupsafe"
version = "3.0.2"
@@ -612,6 +634,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/0d/80/0985960e4b89922cb5a0bac0ed39c5b96cbc1a536a99f30e8c220a996ed9/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:131a3c7689c85f5ad20f9f6fb1b866f402c445b220c19fe4308c0b147ccd2ad9", size = 24098, upload-time = "2024-10-18T15:21:40.813Z" },
]
+[[package]]
+name = "mdit-py-plugins"
+version = "0.5.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "markdown-it-py", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl", hash = "sha256:07a08422fc1936a5d26d146759e9155ea466e842f5ab2f7d2266dd084c8dab1f", size = 57205, upload-time = "2025-08-11T07:25:47.597Z" },
+]
+
[[package]]
name = "mdurl"
version = "0.1.2"
@@ -774,6 +808,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/78/f9/690a8600b93c332de3ab4a344a4ac34f00c8f104917061f779db6a918ed6/pathlib-1.0.1-py3-none-any.whl", hash = "sha256:f35f95ab8b0f59e6d354090350b44a80a80635d22efdedfa84c7ad1cf0a74147", size = 14363, upload-time = "2022-05-04T13:37:20.585Z" },
]
+[[package]]
+name = "platformdirs"
+version = "4.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/23/e8/21db9c9987b0e728855bd57bff6984f67952bea55d6f75e055c46b5383e8/platformdirs-4.4.0.tar.gz", hash = "sha256:ca753cf4d81dc309bc67b0ea38fd15dc97bc30ce419a7f58d13eb3bf14c4febf", size = 21634, upload-time = "2025-08-26T14:32:04.268Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/40/4b/2028861e724d3bd36227adfa20d3fd24c3fc6d52032f4a93c133be5d17ce/platformdirs-4.4.0-py3-none-any.whl", hash = "sha256:abd01743f24e5287cd7a5db3752faf1a2d65353f38ec26d98e25a6db65958c85", size = 18654, upload-time = "2025-08-26T14:32:02.735Z" },
+]
+
[[package]]
name = "pluggy"
version = "1.6.0"
@@ -1122,6 +1165,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ce/fd/901cfa59aaa5b30a99e16876f11abe38b59a1a2c51ffb3d7142bb6089069/starlette-0.47.3-py3-none-any.whl", hash = "sha256:89c0778ca62a76b826101e7c709e70680a1699ca7da6b44d38eb0a7e61fe4b51", size = 72991, upload-time = "2025-08-24T13:36:40.887Z" },
]
+[[package]]
+name = "textual"
+version = "5.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "markdown-it-py", extra = ["linkify", "plugins"], marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "platformdirs", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "pygments", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "rich", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ba/ce/f0f938d33d9bebbf8629e0020be00c560ddfa90a23ebe727c2e5aa3f30cf/textual-5.3.0.tar.gz", hash = "sha256:1b6128b339adef2e298cc23ab4777180443240ece5c232f29b22960efd658d4d", size = 1557651, upload-time = "2025-08-07T12:36:50.342Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/00/2f/f7c8a533bee50fbf5bb37ffc1621e7b2cdd8c9a6301fc51faa35fa50b09d/textual-5.3.0-py3-none-any.whl", hash = "sha256:02a6abc065514c4e21f94e79aaecea1f78a28a85d11d7bfc64abf3392d399890", size = 702671, upload-time = "2025-08-07T12:36:48.272Z" },
+]
+
[[package]]
name = "tokenizers"
version = "0.21.4"
@@ -1217,6 +1276,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/17/69/cd203477f944c353c31bade965f880aa1061fd6bf05ded0726ca845b6ff7/typing_inspection-0.4.1-py3-none-any.whl", hash = "sha256:389055682238f53b04f7badcb49b989835495a96700ced5dab2d8feae4b26f51", size = 14552, upload-time = "2025-05-21T18:55:22.152Z" },
]
+[[package]]
+name = "uc-micro-py"
+version = "1.0.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/91/7a/146a99696aee0609e3712f2b44c6274566bc368dfe8375191278045186b8/uc-micro-py-1.0.3.tar.gz", hash = "sha256:d321b92cff673ec58027c04015fcaa8bb1e005478643ff4a500882eaab88c48a", size = 6043, upload-time = "2024-02-09T16:52:01.654Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/37/87/1f677586e8ac487e29672e4b17455758fce261de06a0d086167bb760361a/uc_micro_py-1.0.3-py3-none-any.whl", hash = "sha256:db1dffff340817673d7b466ec86114a9dc0e9d4d9b5ba229d9d60e5c12600cd5", size = 6229, upload-time = "2024-02-09T16:52:00.371Z" },
+]
+
[[package]]
name = "urllib3"
version = "2.5.0"
← 84c90a6d feat: mlx memory cache for faster ttft
·
back to Exo
·
Prompt length a33787f5 →