electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: model loading progress bars with layer-by-layer tracking

Browse Source 
Emit RunnerLoading progress (layers_loaded/total_layers) during pipeline
parallel model loading and display progress bars in both onboarding step 8
and the main instances view.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Alex Cheema
2026-02-21 11:07:46 -08:00
co-authored by Claude Opus 4.6
parent 885a84c4d1
commit ebf3cdfd72
6 changed files with 118 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dashboard/src/routes/+page.svelte
+94 -21
View File
@@ -701,6 +701,22 @@
return null;
});
// Helper to get onboarding model loading progress (layers loaded)
const onboardingLoadProgress = $derived.by(() => {
if (instanceCount === 0 || !onboardingModelId) return null;
let layersLoaded = 0, totalLayers = 0;
for (const [, inst] of Object.entries(instanceData)) {
if (getInstanceModelId(inst) !== onboardingModelId) continue;
const status = deriveInstanceStatus(inst);
if (status.statusText === "LOADING" && status.totalLayers && status.totalLayers > 0) {
layersLoaded += status.layersLoaded ?? 0;
totalLayers += status.totalLayers;
}
}
if (totalLayers === 0) return null;
return { layersLoaded, totalLayers, percentage: (layersLoaded / totalLayers) * 100 };
});
// Instance launch state
let models = $state<
Array<{
@@ -1720,6 +1736,8 @@
function deriveInstanceStatus(instanceWrapped: unknown): {
statusText: string;
statusClass: string;
layersLoaded?: number;
totalLayers?: number;
} {
const [, instance] = getTagged(instanceWrapped);
if (!instance || typeof instance !== "object") {
@@ -1759,8 +1777,20 @@
if (has("Failed")) return { statusText: "FAILED", statusClass: "failed" };
if (has("Shutdown"))
return { statusText: "SHUTDOWN", statusClass: "inactive" };
if (has("Loading"))
return { statusText: "LOADING", statusClass: "starting" };
if (has("Loading")) {
let layersLoaded = 0, totalLayers = 0;
for (const rid of runnerIds) {
const r = runnersData[rid];
if (!r) continue;
const [kind, payload] = getTagged(r);
if (kind === "RunnerLoading" && payload && typeof payload === "object") {
const p = payload as { layersLoaded?: number; totalLayers?: number };
layersLoaded += p.layersLoaded ?? 0;
totalLayers += p.totalLayers ?? 0;
}
}
return { statusText: "LOADING", statusClass: "starting", layersLoaded, totalLayers };
}
if (has("WarmingUp"))
return { statusText: "WARMING UP", statusClass: "starting" };
if (has("Running"))
@@ -3502,13 +3532,26 @@
</svg>
</div>
<div class="flex justify-center mb-4">
<div
class="w-8 h-8 border-2 border-exo-yellow/15 border-t-exo-yellow/70 rounded-full animate-spin"
></div>
</div>
<p class="text-sm text-white/30 font-sans">Almost ready...</p>
{#if onboardingLoadProgress}
<div class="w-full max-w-xs mx-auto space-y-3">
<div class="relative h-2 bg-white/10 rounded-full overflow-hidden">
<div
class="absolute inset-y-0 left-0 bg-gradient-to-r from-exo-yellow to-exo-yellow-darker rounded-full transition-all duration-500"
style="width: {onboardingLoadProgress.percentage}%"
></div>
</div>
<p class="text-xs text-white/40 font-mono text-center">
{onboardingLoadProgress.layersLoaded} / {onboardingLoadProgress.totalLayers} layers loaded
</p>
</div>
{:else}
<div class="flex justify-center mb-4">
<div
class="w-8 h-8 border-2 border-exo-yellow/15 border-t-exo-yellow/70 rounded-full animate-spin"
></div>
</div>
<p class="text-sm text-white/30 font-sans">Loading...</p>
{/if}
</div>
{:else if onboardingStep === 9}
<!-- Step 9: Ready — centered input with suggestion chips -->
@@ -4397,12 +4440,27 @@
{downloadInfo.statusText}
</div>
{#if isLoading}
<p
class="text-[11px] text-white/50 leading-relaxed"
>
Loading model into memory for fast
inference...
</p>
{@const loadStatus = deriveInstanceStatus(instance)}
{#if loadStatus.totalLayers && loadStatus.totalLayers > 0}
<div class="mt-1 space-y-1">
<div class="flex justify-between text-xs font-mono">
<span class="text-yellow-400">{((loadStatus.layersLoaded ?? 0) / loadStatus.totalLayers * 100).toFixed(0)}%</span>
<span class="text-exo-light-gray">{loadStatus.layersLoaded ?? 0} / {loadStatus.totalLayers} layers</span>
</div>
<div class="relative h-1.5 bg-exo-black/60 rounded-sm overflow-hidden">
<div
class="absolute inset-y-0 left-0 bg-gradient-to-r from-yellow-500 to-yellow-400 transition-all duration-300"
style="width: {(loadStatus.layersLoaded ?? 0) / loadStatus.totalLayers * 100}%"
></div>
</div>
</div>
{:else}
<p
class="text-[11px] text-white/50 leading-relaxed"
>
Loading model into memory...
</p>
{/if}
{:else if isReady || isRunning}
<p
class="text-[11px] text-green-400/70 leading-relaxed"
@@ -5301,12 +5359,27 @@
{downloadInfo.statusText}
</div>
{#if isLoading}
<p
class="text-[11px] text-white/50 leading-relaxed"
>
Loading model into memory for fast
inference...
</p>
{@const loadStatus = deriveInstanceStatus(instance)}
{#if loadStatus.totalLayers && loadStatus.totalLayers > 0}
<div class="mt-1 space-y-1">
<div class="flex justify-between text-xs font-mono">
<span class="text-yellow-400">{((loadStatus.layersLoaded ?? 0) / loadStatus.totalLayers * 100).toFixed(0)}%</span>
<span class="text-exo-light-gray">{loadStatus.layersLoaded ?? 0} / {loadStatus.totalLayers} layers</span>
</div>
<div class="relative h-1.5 bg-exo-black/60 rounded-sm overflow-hidden">
<div
class="absolute inset-y-0 left-0 bg-gradient-to-r from-yellow-500 to-yellow-400 transition-all duration-300"
style="width: {(loadStatus.layersLoaded ?? 0) / loadStatus.totalLayers * 100}%"
></div>
</div>
</div>
{:else}
<p
class="text-[11px] text-white/50 leading-relaxed"
>
Loading model into memory...
</p>
{/if}
{:else if isReady || isRunning}
<p
class="text-[11px] text-green-400/70 leading-relaxed"
src/exo/shared/types/worker/runners.py
+2 -1
View File
@@ -34,7 +34,8 @@ class RunnerConnected(BaseRunnerStatus):
class RunnerLoading(BaseRunnerStatus):
pass
layers_loaded: int = 0
total_layers: int = 0
class RunnerLoaded(BaseRunnerStatus):
src/exo/worker/engines/mlx/auto_parallel.py
+6 -1
View File
@@ -47,6 +47,7 @@ if TYPE_CHECKING:
from mlx_lm.models.cache import Cache
TimeoutCallback = Callable[[], None]
LayerLoadedCallback = Callable[[int, int], None] # (layers_loaded, total_layers)
def eval_with_timeout(
@@ -221,6 +222,7 @@ def pipeline_auto_parallel(
model: nn.Module,
group: mx.distributed.Group,
model_shard_meta: PipelineShardMetadata,
on_layer_loaded: LayerLoadedCallback | None = None,
) -> nn.Module:
"""
Automatically parallelize a model across multiple devices.
@@ -238,8 +240,11 @@ def pipeline_auto_parallel(
device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
layers = layers[start_layer:end_layer]
for layer in layers:
total = len(layers)
for i, layer in enumerate(layers):
mx.eval(layer) # type: ignore
if on_layer_loaded is not None:
on_layer_loaded(i + 1, total)
layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
layers[-1] = PipelineLastLayer(
src/exo/worker/engines/mlx/utils_mlx.py
+5 -2
View File
@@ -55,6 +55,7 @@ from exo.shared.types.worker.shards import (
)
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.auto_parallel import (
LayerLoadedCallback,
TimeoutCallback,
eval_with_timeout,
pipeline_auto_parallel,
@@ -172,6 +173,7 @@ def load_mlx_items(
bound_instance: BoundInstance,
group: Group | None,
on_timeout: TimeoutCallback | None = None,
on_layer_loaded: LayerLoadedCallback | None = None,
) -> tuple[Model, TokenizerWrapper]:
if group is None:
logger.info(f"Single device used for {bound_instance.instance}")
@@ -186,7 +188,7 @@ def load_mlx_items(
logger.info("Starting distributed init")
start_time = time.perf_counter()
model, tokenizer = shard_and_load(
bound_instance.bound_shard, group=group, on_timeout=on_timeout
bound_instance.bound_shard, group=group, on_timeout=on_timeout, on_layer_loaded=on_layer_loaded
)
end_time = time.perf_counter()
logger.info(
@@ -202,6 +204,7 @@ def shard_and_load(
shard_metadata: ShardMetadata,
group: Group,
on_timeout: TimeoutCallback | None = None,
on_layer_loaded: LayerLoadedCallback | None = None,
) -> tuple[nn.Module, TokenizerWrapper]:
model_path = build_model_path(shard_metadata.model_card.model_id)
@@ -245,7 +248,7 @@ def shard_and_load(
model = tensor_auto_parallel(model, group, timeout_seconds, on_timeout)
case PipelineShardMetadata():
logger.info(f"loading model from {model_path} with pipeline parallelism")
model = pipeline_auto_parallel(model, group, shard_metadata)
model = pipeline_auto_parallel(model, group, shard_metadata, on_layer_loaded=on_layer_loaded)
eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
case CfgShardMetadata():
raise ValueError(
src/exo/worker/runner/llm_inference/runner.py
+10 -2
View File
@@ -150,7 +150,8 @@ def main(
case LoadModel() if (
isinstance(current_status, RunnerConnected) and group is not None
) or (isinstance(current_status, RunnerIdle) and group is None):
current_status = RunnerLoading()
total_layers = shard_metadata.end_layer - shard_metadata.start_layer
current_status = RunnerLoading(layers_loaded=0, total_layers=total_layers)
logger.info("runner loading")
event_sender.send(
RunnerStatusUpdated(
@@ -170,11 +171,18 @@ def main(
)
time.sleep(0.5)
def on_layer_loaded(layers_loaded: int, total: int) -> None:
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=RunnerLoading(layers_loaded=layers_loaded, total_layers=total)
)
)
assert (
ModelTask.TextGeneration in shard_metadata.model_card.tasks
), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
inference_model, tokenizer = load_mlx_items(
bound_instance, group, on_timeout=on_model_load_timeout
bound_instance, group, on_timeout=on_model_load_timeout, on_layer_loaded=on_layer_loaded
)
logger.info(
f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}"
src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+1 -1
View File
@@ -224,7 +224,7 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerConnected()),
TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Running),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoading()),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoading(layers_loaded=0, total_layers=32)),
TaskAcknowledged(task_id=LOAD_TASK_ID),
TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Complete),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoaded()),