From 17f9b583a48def54c99d35a5815a4fdc3e81f090 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Sat, 3 Jan 2026 20:01:49 +0000 Subject: [PATCH 1/2] Task Deduplication (#1062) --- src/exo/shared/types/worker/runners.py | 5 +++ src/exo/worker/plan.py | 6 +++ src/exo/worker/runner/runner.py | 13 ++++--- src/exo/worker/runner/runner_supervisor.py | 37 +++++++++++++++++-- src/exo/worker/tests/unittests/conftest.py | 7 ++-- .../test_runner/test_event_ordering.py | 4 ++ 6 files changed, 58 insertions(+), 14 deletions(-) diff --git a/src/exo/shared/types/worker/runners.py b/src/exo/shared/types/worker/runners.py index 402f9619..ff5026e3 100644 --- a/src/exo/shared/types/worker/runners.py +++ b/src/exo/shared/types/worker/runners.py @@ -53,6 +53,10 @@ class RunnerRunning(BaseRunnerStatus): pass +class RunnerShuttingDown(BaseRunnerStatus): + pass + + class RunnerShutdown(BaseRunnerStatus): pass @@ -70,6 +74,7 @@ RunnerStatus = ( | RunnerWarmingUp | RunnerReady | RunnerRunning + | RunnerShuttingDown | RunnerShutdown | RunnerFailed ) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index d72f2e2d..70b2e02d 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -274,6 +274,12 @@ def _pending_tasks( if task.instance_id != runner.bound_instance.instance.instance_id: continue + # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough + # however, realistically the task status should be set to completed by the LAST runner, so this is a true race + # the actual solution is somewhat deeper than this bypass - TODO! + if task.task_id in runner.completed: + continue + # TODO: Check ordering aligns with MLX distributeds expectations. if isinstance(runner.status, RunnerReady) and all( diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index b3e8a506..14510c1f 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -32,6 +32,7 @@ from exo.shared.types.worker.runners import ( RunnerReady, RunnerRunning, RunnerShutdown, + RunnerShuttingDown, RunnerStatus, RunnerWarmingUp, ) @@ -187,13 +188,14 @@ def main( current_status = RunnerReady() logger.info("runner ready") case Shutdown(): + current_status = RunnerShuttingDown() logger.info("runner shutting down") event_sender.send( - TaskStatusUpdated( - task_id=task.task_id, task_status=TaskStatus.Complete + RunnerStatusUpdated( + runner_id=runner_id, runner_status=current_status ) ) - break + current_status = RunnerShutdown() case _: raise ValueError( f"Received {task.__class__.__name__} outside of state machine in {current_status=}" @@ -208,9 +210,8 @@ def main( runner_id=runner_id, runner_status=current_status ) ) - event_sender.send( - RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerShutdown()) - ) + if isinstance(current_status, RunnerShutdown): + break except ClosedResourceError: logger.warning("runner communication closed unexpectedly") except Exception as e: diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 3c916a8b..fc17cddc 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -14,13 +14,23 @@ from anyio import ( from anyio.abc import TaskGroup from loguru import logger -from exo.shared.types.events import Event, RunnerStatusUpdated, TaskAcknowledged -from exo.shared.types.tasks import Task, TaskId +from exo.shared.types.events import ( + Event, + RunnerStatusUpdated, + TaskAcknowledged, + TaskStatusUpdated, +) +from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import ( + RunnerConnecting, RunnerFailed, RunnerIdle, + RunnerLoading, + RunnerRunning, + RunnerShuttingDown, RunnerStatus, + RunnerWarmingUp, ) from exo.shared.types.worker.shards import ShardMetadata from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel @@ -39,10 +49,10 @@ class RunnerSupervisor: _ev_recv: MpReceiver[Event] _task_sender: MpSender[Task] _event_sender: Sender[Event] - # err_path: str _tg: TaskGroup | None = field(default=None, init=False) status: RunnerStatus = field(default_factory=RunnerIdle, init=False) pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False) + completed: set[TaskId] = field(default_factory=set, init=False) @classmethod def create( @@ -77,7 +87,6 @@ class RunnerSupervisor: _ev_recv=ev_recv, _task_sender=task_sender, _event_sender=event_sender, - # err_path=err_path, ) return self @@ -118,6 +127,10 @@ class RunnerSupervisor: self._tg.cancel_scope.cancel() async def start_task(self, task: Task): + if task.task_id in self.completed: + logger.info( + f"Skipping invalid task {task} as it has already been completed" + ) logger.info(f"Starting task {task}") event = anyio.Event() self.pending[task.task_id] = event @@ -138,6 +151,22 @@ class RunnerSupervisor: if isinstance(event, TaskAcknowledged): self.pending.pop(event.task_id).set() continue + if ( + isinstance(event, TaskStatusUpdated) + and event.task_status == TaskStatus.Complete + ): + # If a task has just been completed, we should be working on it. + assert isinstance( + self.status, + ( + RunnerRunning, + RunnerWarmingUp, + RunnerLoading, + RunnerConnecting, + RunnerShuttingDown, + ), + ) + self.completed.add(event.task_id) await self._event_sender.send(event) except (ClosedResourceError, BrokenResourceError) as e: await self._check_runner(e) diff --git a/src/exo/worker/tests/unittests/conftest.py b/src/exo/worker/tests/unittests/conftest.py index 18f2e0de..dea085fe 100644 --- a/src/exo/worker/tests/unittests/conftest.py +++ b/src/exo/worker/tests/unittests/conftest.py @@ -1,11 +1,9 @@ -from __future__ import annotations - -from dataclasses import dataclass +from dataclasses import dataclass, field from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.models import ModelId, ModelMetadata -from exo.shared.types.tasks import BaseTask +from exo.shared.types.tasks import BaseTask, TaskId from exo.shared.types.worker.instances import ( BoundInstance, Instance, @@ -21,6 +19,7 @@ from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata class FakeRunnerSupervisor: bound_instance: BoundInstance status: RunnerStatus + completed: set[TaskId] = field(default_factory=set) class OtherTask(BaseTask): diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py index abe9c603..954052c3 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py +++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py @@ -34,6 +34,7 @@ from exo.shared.types.worker.runners import ( RunnerReady, RunnerRunning, RunnerShutdown, + RunnerShuttingDown, RunnerWarmingUp, ) from exo.utils.channels import mp_channel @@ -199,6 +200,9 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch): RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerReady()), TaskStatusUpdated(task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Running), TaskAcknowledged(task_id=SHUTDOWN_TASK_ID), + RunnerStatusUpdated( + runner_id=RUNNER_1_ID, runner_status=RunnerShuttingDown() + ), TaskStatusUpdated( task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete ), From 47b8e0ce12e3bd0b417a13d50188d8fc2e9dc73a Mon Sep 17 00:00:00 2001 From: Drifter4242 Date: Mon, 5 Jan 2026 11:27:14 +0000 Subject: [PATCH 2/2] feat: remember last launch settings (model, sharding, instance type) (#1028) ## Motivation Saves the last launch settings, so that the next time you run exo it will default to the same launch settings. This is just a small quality of life improvement. ## Changes When you launch it saves the settings to the web browser local storage. When it fills out the model list, it reads the settings and sets the default. I reviewed, tested and edited the code, but some of the code was written by Claude Opus. I hope that's ok. ## Why It Works See above ## Test Plan ### Manual Testing I have two Macbook Studio M3 Ultras, each with 512Gb ram, connected with Thunderbolt 5. I ran Kimi K2 Thinking with MLX Ring and Tensor Split. I ran exo multiple times to confirm that the default works. ### Automated Testing No changes to automated testing. --- dashboard/src/routes/+page.svelte | 67 +++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 4 deletions(-) diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 9f75eac8..abc603ff 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -51,6 +51,59 @@ const sidebarVisible = $derived(chatSidebarVisible()); let selectedSharding = $state<'Pipeline' | 'Tensor'>('Pipeline'); type InstanceMeta = 'MlxRing' | 'MlxIbv' | 'MlxJaccl'; + // Launch defaults persistence + const LAUNCH_DEFAULTS_KEY = 'exo-launch-defaults'; + interface LaunchDefaults { + modelId: string | null; + sharding: 'Pipeline' | 'Tensor'; + instanceType: InstanceMeta; + minNodes: number; + } + + function saveLaunchDefaults(): void { + const defaults: LaunchDefaults = { + modelId: selectedPreviewModelId(), + sharding: selectedSharding, + instanceType: selectedInstanceType, + minNodes: selectedMinNodes, + }; + try { + localStorage.setItem(LAUNCH_DEFAULTS_KEY, JSON.stringify(defaults)); + } catch (e) { + console.warn('Failed to save launch defaults:', e); + } + } + + function loadLaunchDefaults(): LaunchDefaults | null { + try { + const stored = localStorage.getItem(LAUNCH_DEFAULTS_KEY); + if (!stored) return null; + return JSON.parse(stored) as LaunchDefaults; + } catch (e) { + console.warn('Failed to load launch defaults:', e); + return null; + } + } + + function applyLaunchDefaults(availableModels: Array<{id: string}>, maxNodes: number): void { + const defaults = loadLaunchDefaults(); + if (!defaults) return; + + // Apply sharding and instance type unconditionally + selectedSharding = defaults.sharding; + selectedInstanceType = defaults.instanceType; + + // Apply minNodes if valid (between 1 and maxNodes) + if (defaults.minNodes && defaults.minNodes >= 1 && defaults.minNodes <= maxNodes) { + selectedMinNodes = defaults.minNodes; + } + + // Only apply model if it exists in the available models + if (defaults.modelId && availableModels.some(m => m.id === defaults.modelId)) { + selectPreviewModel(defaults.modelId); + } + } + let selectedInstanceType = $state('MlxRing'); let selectedMinNodes = $state(1); let minNodesInitialized = $state(false); @@ -298,6 +351,9 @@ function toggleInstanceDownloadDetails(nodeId: string): void { const data = await response.json(); // API returns { data: [{ id, name }] } format models = data.data || []; + // Restore last launch defaults if available + const currentNodeCount = topologyData() ? Object.keys(topologyData()!.nodes).length : 1; + applyLaunchDefaults(models, currentNodeCount); } } catch (error) { console.error('Failed to fetch models:', error); @@ -988,6 +1044,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void { function handleSliderMouseUp() { isDraggingSlider = false; + saveLaunchDefaults(); } // Handle touch events for mobile @@ -1007,6 +1064,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void { function handleSliderTouchEnd() { isDraggingSlider = false; + saveLaunchDefaults(); } const nodeCount = $derived(data ? Object.keys(data.nodes).length : 0); @@ -1464,6 +1522,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void { onclick={() => { if (modelCanFit) { selectPreviewModel(model.id); + saveLaunchDefaults(); isModelDropdownOpen = false; modelDropdownSearch = ''; } @@ -1497,7 +1556,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
Sharding: