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

Implement storage management

Browse Source
This commit is contained in:
ciaranbor
2026-04-01 19:14:11 +01:00
parent 94a7264cbf
commit f87d705b9c
15 changed files with 1821 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dashboard/src/lib/stores/app.svelte.ts
+109
View File
@@ -8,6 +8,11 @@
*/
import { browser } from "$app/environment";
import { addToast } from "$lib/stores/toast.svelte";
import {
getDownloadTag,
extractModelIdFromDownload,
} from "$lib/utils/downloads";
// UUID generation fallback for browsers without crypto.randomUUID
function generateUUID(): string {
@@ -254,6 +259,13 @@ interface RawStateResponse {
string,
{ total: { inBytes: number }; available: { inBytes: number } }
>;
nodeStorageConfig?: Record<
string,
{
maxStorage: { inBytes: number } | null;
storagePolicy: "manual" | "auto-evict";
}
>;
}
export interface MessageAttachment {
@@ -568,6 +580,15 @@ class AppStore {
>
>({});
nodeRdmaCtl = $state<Record<string, { enabled: boolean }>>({});
nodeStorageConfig = $state<
Record<
string,
{
maxStorage: { inBytes: number } | null;
storagePolicy: "manual" | "auto-evict";
}
>
>({});
nodeThunderboltBridge = $state<
Record<
string,
@@ -1311,6 +1332,7 @@ class AppStore {
this.runners = data.runners;
}
if (data.downloads) {
this.detectEvictions(this.downloads, data.downloads);
this.downloads = data.downloads;
}
if (data.nodeDisk) {
@@ -1326,6 +1348,7 @@ class AppStore {
this.thunderboltBridgeCycles = data.thunderboltBridgeCycles ?? [];
// Thunderbolt bridge status per node
this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
this.nodeStorageConfig = data.nodeStorageConfig ?? {};
this.lastUpdate = Date.now();
// Connection recovered
if (!this.isConnected) {
@@ -1438,6 +1461,63 @@ class AppStore {
}
}
/**
* Detect models that were evicted (completed -> not completed) and show toasts.
*/
private detectEvictions(
oldDownloads: Record<string, unknown[]>,
newDownloads: Record<string, unknown[]>,
) {
if (Object.keys(oldDownloads).length === 0) return;
const wasCompleted = new Map<string, Set<string>>();
for (const [nodeId, entries] of Object.entries(oldDownloads)) {
if (!Array.isArray(entries)) continue;
for (const entry of entries) {
const tagged = getDownloadTag(entry);
if (!tagged || tagged[0] !== "DownloadCompleted") continue;
const modelId = extractModelIdFromDownload(tagged[1]);
if (!modelId) continue;
if (!wasCompleted.has(nodeId)) wasCompleted.set(nodeId, new Set());
wasCompleted.get(nodeId)!.add(modelId);
}
}
for (const [nodeId, modelIds] of wasCompleted) {
const newEntries = newDownloads[nodeId];
const stillCompleted = new Set<string>();
if (Array.isArray(newEntries)) {
for (const entry of newEntries) {
const tagged = getDownloadTag(entry);
if (!tagged || tagged[0] !== "DownloadCompleted") continue;
const mid = extractModelIdFromDownload(tagged[1]);
if (mid) stillCompleted.add(mid);
}
}
for (const modelId of modelIds) {
if (!stillCompleted.has(modelId)) {
const shortName = modelId.split("/").pop() ?? modelId;
const nodeLabel = this.getNodeFriendlyName(nodeId);
addToast({
type: "info",
message: `Model evicted on ${nodeLabel}: ${shortName}`,
duration: 6000,
});
}
}
}
}
private getNodeFriendlyName(nodeId: string): string {
const nodeInfo = this.topologyData?.nodes?.[nodeId];
return (
nodeInfo?.friendly_name ??
nodeInfo?.system_info?.chip ??
nodeId.slice(0, 4)
);
}
/**
* Handle topology changes - clean up filter and re-fetch if needed
*/
@@ -3279,6 +3359,29 @@ class AppStore {
}
}
async setStorageConfig(
nodeId: string,
maxStorageGb: number | null,
storagePolicy: "manual" | "auto-evict",
): Promise<void> {
try {
const response = await fetch(`/storage/${encodeURIComponent(nodeId)}`, {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ maxStorageGb, storagePolicy }),
});
if (!response.ok) {
const errorText = await response.text();
throw new Error(
`Failed to set storage config: ${response.status} - ${errorText}`,
);
}
} catch (error) {
console.error("Error setting storage config:", error);
throw error;
}
}
/**
* List all available traces
*/
@@ -3356,6 +3459,7 @@ export const instances = () => appStore.instances;
export const runners = () => appStore.runners;
export const downloads = () => appStore.downloads;
export const nodeDisk = () => appStore.nodeDisk;
export const nodeStorageConfig = () => appStore.nodeStorageConfig;
export const placementPreviews = () => appStore.placementPreviews;
export const selectedPreviewModelId = () => appStore.selectedPreviewModelId;
export const isLoadingPreviews = () => appStore.isLoadingPreviews;
@@ -3479,6 +3583,11 @@ export const startDownload = (nodeId: string, shardMetadata: object) =>
appStore.startDownload(nodeId, shardMetadata);
export const deleteDownload = (nodeId: string, modelId: string) =>
appStore.deleteDownload(nodeId, modelId);
export const setStorageConfig = (
nodeId: string,
maxStorageGb: number | null,
storagePolicy: "manual" | "auto-evict",
) => appStore.setStorageConfig(nodeId, maxStorageGb, storagePolicy);
// Trace actions
export const listTraces = () => appStore.listTraces();
dashboard/src/routes/+page.svelte
+68 -1
View File
@@ -1581,12 +1581,14 @@
progress: DownloadProgress | null;
perNode: NodeDownloadStatus[];
failedError: string | null;
rejectedError: string | null;
} {
const empty = {
isDownloading: false,
progress: null,
perNode: [] as NodeDownloadStatus[],
failedError: null,
rejectedError: null,
};
if (!downloadsData || Object.keys(downloadsData).length === 0) {
@@ -1628,6 +1630,19 @@
(downloadPayload.errorMessage as string) ||
(downloadPayload.error_message as string) ||
"Download failed",
rejectedError: null,
};
}
// DownloadRejected — storage limit exceeded
if (downloadKind === "DownloadRejected") {
return {
isDownloading: false,
progress: null,
perNode: Array.from(perNodeMap.values()),
failedError: null,
rejectedError:
(downloadPayload.reason as string) || "Storage limit exceeded",
};
}
@@ -1722,6 +1737,7 @@
progress: null,
perNode,
failedError: null,
rejectedError: null,
};
}
@@ -1742,6 +1758,7 @@
},
perNode,
failedError: null,
rejectedError: null,
};
}
@@ -1826,6 +1843,17 @@
};
}
if (result.rejectedError) {
return {
isDownloading: false,
isFailed: true,
errorMessage: result.rejectedError,
progress: null,
statusText: "REJECTED",
perNode: [],
};
}
if (!result.isDownloading) {
const statusInfo = deriveInstanceStatus(instanceWrapped);
return {
@@ -2461,6 +2489,7 @@
// ── Instance status transition toasts ──
// Track previous statuses so we can detect meaningful transitions and fire toasts.
let previousInstanceStatuses: Record<string, string> = {};
let previousInstanceModelIds: Record<string, string> = {};
$effect(() => {
const currentStatuses: Record<string, string> = {};
@@ -2475,7 +2504,13 @@
if (Object.keys(prev).length > 0) {
for (const [id, currentStatus] of Object.entries(currentStatuses)) {
const prevStatus = prev[id];
if (!prevStatus || prevStatus === currentStatus) continue;
if (prevStatus === currentStatus) continue;
if (
!prevStatus &&
currentStatus !== "REJECTED" &&
currentStatus !== "FAILED"
)
continue;
const modelId = getInstanceModelId(instanceData[id]);
const shortName = modelId
@@ -2509,6 +2544,14 @@
addToast({ type: "error", message: `Model failed: ${shortName}` });
}
if (prevStatus !== "REJECTED" && currentStatus === "REJECTED") {
addToast({
type: "warning",
message: `Storage limit exceeded: ${shortName}`,
duration: 8000,
});
}
// Any -> Shutdown
if (prevStatus !== "SHUTDOWN" && currentStatus === "SHUTDOWN") {
addToast({ type: "info", message: `Model shut down: ${shortName}` });
@@ -2516,7 +2559,31 @@
}
}
// Detect instances that disappeared while in early states (e.g. rejected download)
if (Object.keys(prev).length > 0) {
for (const [id, prevStatus] of Object.entries(prev)) {
if (id in currentStatuses) continue; // still exists
if (prevStatus === "PREPARING" || prevStatus === "DOWNLOADING") {
const modelId = previousInstanceModelIds[id];
const shortName = modelId
? (modelId.split("/").pop() ?? modelId)
: id.slice(0, 8);
addToast({
type: "warning",
message: `Download cancelled: ${shortName} — insufficient storage`,
duration: 8000,
});
}
}
}
previousInstanceStatuses = currentStatuses;
const modelIds: Record<string, string> = {};
for (const [id, inst] of Object.entries(instanceData)) {
const mid = getInstanceModelId(inst);
if (mid) modelIds[id] = mid;
}
previousInstanceModelIds = modelIds;
});
// ── Connection status toasts ──
dashboard/src/routes/downloads/+page.svelte
+321 -12
View File
@@ -6,16 +6,19 @@
topologyData,
downloads,
nodeDisk,
nodeStorageConfig,
refreshState,
lastUpdate as lastUpdateStore,
startDownload,
deleteDownload,
setStorageConfig,
} from "$lib/stores/app.svelte";
import {
getDownloadTag,
extractModelIdFromDownload,
extractShardMetadata,
} from "$lib/utils/downloads";
import { addToast } from "$lib/stores/toast.svelte";
import HeaderNav from "$lib/components/HeaderNav.svelte";
type CellStatus =
@@ -36,6 +39,14 @@
modelDirectory?: string;
}
| { kind: "failed"; modelDirectory?: string }
| {
kind: "rejected";
reason: string;
requiredBytes: number;
availableBytes: number;
limitBytes: number;
modelDirectory?: string;
}
| { kind: "not_present" };
type ModelCardInfo = {
@@ -61,11 +72,14 @@
label: string;
diskAvailable?: number;
diskTotal?: number;
storageLimit?: number;
storagePolicy?: "manual" | "auto-evict";
};
const data = $derived(topologyData());
const downloadsData = $derived(downloads());
const nodeDiskData = $derived(nodeDisk());
const storageConfigData = $derived(nodeStorageConfig());
function getNodeLabel(nodeId: string): string {
const node = data?.nodes?.[nodeId];
@@ -122,10 +136,37 @@
return Math.min(100, Math.max(0, value as number));
}
function getNodeUsedStorage(nodeId: string): number {
const nodeDownloads = downloadsData?.[nodeId];
if (!nodeDownloads || !Array.isArray(nodeDownloads)) return 0;
let total = 0;
for (const entry of nodeDownloads) {
const tagged = getDownloadTag(entry);
if (!tagged) continue;
const [tag, payload] = tagged;
if (tag === "DownloadCompleted") {
total += getBytes(payload.total);
} else if (tag === "DownloadOngoing") {
const prog = (payload.download_progress ?? payload.downloadProgress) as
| Record<string, unknown>
| undefined;
if (prog) total += getBytes(prog.downloaded);
}
}
return total;
}
function storageBarColor(percent: number): string {
if (percent >= 90) return "bg-red-500";
if (percent >= 70) return "bg-yellow-500";
return "bg-green-500";
}
const CELL_PRIORITY: Record<CellStatus["kind"], number> = {
completed: 4,
downloading: 3,
pending: 2,
completed: 5,
downloading: 4,
pending: 3,
rejected: 2,
failed: 1,
not_present: 0,
};
@@ -178,6 +219,43 @@
let nodeColumns = $state<NodeColumn[]>([]);
let infoRow = $state<ModelRow | null>(null);
// Storage config modal state
let storageConfigNode = $state<NodeColumn | null>(null);
let configMaxGb = $state<number | null>(null);
let configNoLimit = $state(true);
let configPolicy = $state<"manual" | "auto-evict">("manual");
let configSaving = $state(false);
function openStorageConfig(col: NodeColumn) {
storageConfigNode = col;
if (col.storageLimit != null) {
configNoLimit = false;
configMaxGb = Math.round(col.storageLimit / 1024 ** 3);
} else {
configNoLimit = true;
configMaxGb = null;
}
configPolicy = col.storagePolicy ?? "manual";
}
async function saveStorageConfig() {
if (!storageConfigNode) return;
configSaving = true;
try {
const maxGb = configNoLimit ? null : configMaxGb;
await setStorageConfig(storageConfigNode.nodeId, maxGb, configPolicy);
storageConfigNode = null;
refreshState();
} catch (error) {
addToast({
type: "error",
message: `Failed to save storage config: ${error instanceof Error ? error.message : String(error)}`,
});
} finally {
configSaving = false;
}
}
$effect(() => {
try {
if (!downloadsData || Object.keys(downloadsData).length === 0) {
@@ -189,11 +267,14 @@
const allNodeIds = Object.keys(downloadsData);
const columns: NodeColumn[] = allNodeIds.map((nodeId) => {
const diskInfo = nodeDiskData?.[nodeId];
const storageConfig = storageConfigData?.[nodeId];
return {
nodeId,
label: getNodeLabel(nodeId),
diskAvailable: diskInfo?.available?.inBytes,
diskTotal: diskInfo?.total?.inBytes,
storageLimit: storageConfig?.maxStorage?.inBytes ?? undefined,
storagePolicy: storageConfig?.storagePolicy,
};
});
@@ -236,7 +317,11 @@
let cell: CellStatus;
if (tag === "DownloadCompleted") {
const totalBytes = getBytes(payload.total);
cell = { kind: "completed", totalBytes, modelDirectory };
cell = {
kind: "completed",
totalBytes,
modelDirectory,
};
} else if (tag === "DownloadOngoing") {
const rawProgress =
payload.download_progress ?? payload.downloadProgress ?? {};
@@ -257,6 +342,15 @@
etaMs,
modelDirectory,
};
} else if (tag === "DownloadRejected") {
cell = {
kind: "rejected",
reason: (payload.reason as string) ?? "Storage limit exceeded",
requiredBytes: getBytes(payload.required),
availableBytes: getBytes(payload.available),
limitBytes: getBytes(payload.limit),
modelDirectory,
};
} else if (tag === "DownloadFailed") {
cell = { kind: "failed", modelDirectory };
} else {
@@ -406,15 +500,52 @@
Model
</th>
{#each nodeColumns as col}
{@const usedStorage = getNodeUsedStorage(col.nodeId)}
{@const storageMax = col.storageLimit ?? col.diskTotal ?? 0}
{@const storagePercent =
storageMax > 0
? Math.min(100, (usedStorage / storageMax) * 100)
: 0}
<th
class="px-4 py-3 text-[11px] uppercase tracking-wider text-exo-light-gray font-medium text-center whitespace-nowrap min-w-[120px]"
>
<div>{col.label}</div>
{#if col.diskAvailable != null}
<div
class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
<div class="flex items-center justify-center gap-1">
<span>{col.label}</span>
<button
type="button"
class="p-0.5 rounded hover:bg-white/10 transition-colors"
onclick={() => openStorageConfig(col)}
title="Storage settings"
aria-label="Storage settings for {col.label}"
>
{formatBytes(col.diskAvailable)} free
<svg
class="w-3.5 h-3.5 text-white/40 hover:text-exo-yellow transition-colors"
viewBox="0 0 20 20"
fill="currentColor"
>
<path
fill-rule="evenodd"
d="M11.49 3.17c-.38-1.56-2.6-1.56-2.98 0a1.532 1.532 0 01-2.286.948c-1.372-.836-2.942.734-2.106 2.106.54.886.061 2.042-.947 2.287-1.561.379-1.561 2.6 0 2.978a1.532 1.532 0 01.947 2.287c-.836 1.372.734 2.942 2.106 2.106a1.532 1.532 0 012.287.947c.379 1.561 2.6 1.561 2.978 0a1.533 1.533 0 012.287-.947c1.372.836 2.942-.734 2.106-2.106a1.533 1.533 0 01.947-2.287c1.561-.379 1.561-2.6 0-2.978a1.532 1.532 0 01-.947-2.287c.836-1.372-.734-2.942-2.106-2.106a1.532 1.532 0 01-2.287-.947zM10 13a3 3 0 100-6 3 3 0 000 6z"
clip-rule="evenodd"
/>
</svg>
</button>
</div>
{#if storageMax > 0}
<div class="text-[9px] normal-case tracking-normal mt-1">
<div
class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
>
<div
class="h-full rounded-full transition-all duration-300 {storageBarColor(
storagePercent,
)}"
style="width: {storagePercent.toFixed(1)}%"
></div>
</div>
<div class="text-white/60 mt-0.5">
{formatBytes(usedStorage)} / {formatBytes(storageMax)}
</div>
</div>
{/if}
</th>
@@ -610,6 +741,52 @@
<span class="text-white/40 text-sm">...</span>
{/if}
</div>
{:else if cell.kind === "rejected"}
<div
class="flex flex-col items-center gap-1"
title={cell.reason}
>
<svg
class="w-7 h-7 text-orange-400"
viewBox="0 0 20 20"
fill="currentColor"
>
<path
fill-rule="evenodd"
d="M8.257 3.099c.765-1.36 2.722-1.36 3.486 0l5.58 9.92c.75 1.334-.213 2.98-1.742 2.98H4.42c-1.53 0-2.493-1.646-1.743-2.98l5.58-9.92zM11 13a1 1 0 11-2 0 1 1 0 012 0zm-1-8a1 1 0 00-1 1v3a1 1 0 002 0V6a1 1 0 00-1-1z"
clip-rule="evenodd"
></path>
</svg>
<span class="text-[10px] text-orange-400/80"
>Need {formatBytes(cell.requiredBytes)}</span
>
<span class="text-[10px] text-white/50"
>{formatBytes(cell.availableBytes)} avail</span
>
{#if row.shardMetadata}
<button
type="button"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Retry download on this node"
>
<svg
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
stroke-width="2"
>
<path
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
stroke-linecap="round"
stroke-linejoin="round"
></path>
</svg>
</button>
{/if}
</div>
{:else if cell.kind === "failed"}
<div
class="flex flex-col items-center gap-1"
@@ -812,9 +989,11 @@
? 'bg-green-500/10 text-green-400/80 border border-green-500/20'
: cellStatus.kind === 'downloading'
? 'bg-exo-yellow/10 text-exo-yellow/80 border border-exo-yellow/20'
: cellStatus.kind === 'failed'
? 'bg-red-500/10 text-red-400/80 border border-red-500/20'
: 'bg-white/5 text-white/50 border border-white/10'}"
: cellStatus.kind === 'rejected'
? 'bg-orange-500/10 text-orange-400/80 border border-orange-500/20'
: cellStatus.kind === 'failed'
? 'bg-red-500/10 text-red-400/80 border border-red-500/20'
: 'bg-white/5 text-white/50 border border-white/10'}"
>
{col.label}
{#if cellStatus.kind === "downloading" && "percentage" in cellStatus}
@@ -839,6 +1018,136 @@
</div>
{/if}
<!-- Storage config modal -->
{#if storageConfigNode}
<div
class="fixed inset-0 z-[60] bg-black/60"
transition:fade={{ duration: 150 }}
onclick={() => (storageConfigNode = null)}
role="presentation"
></div>
<div
class="fixed z-[60] top-1/2 left-1/2 -translate-x-1/2 -translate-y-1/2 w-[min(80vw,360px)] bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-2xl p-4"
transition:fly={{ y: 10, duration: 200, easing: cubicOut }}
role="dialog"
aria-modal="true"
onkeydown={(e) => {
if (e.key === "Escape") storageConfigNode = null;
}}
>
<div class="flex items-start justify-between mb-4">
<h3 class="font-mono text-sm text-white">
Storage — {storageConfigNode.label}
</h3>
<button
type="button"
class="p-1 rounded hover:bg-white/10 transition-colors text-white/50"
onclick={() => (storageConfigNode = null)}
aria-label="Close storage settings"
>
<svg class="w-4 h-4" viewBox="0 0 24 24" fill="currentColor">
<path
d="M19 6.41L17.59 5 12 10.59 6.41 5 5 6.41 10.59 12 5 17.59 6.41 19 12 13.41 17.59 19 19 17.59 13.41 12 19 6.41z"
/>
</svg>
</button>
</div>
<div class="space-y-4">
<!-- No limit checkbox -->
<label class="flex items-center gap-2 cursor-pointer">
<input
type="checkbox"
bind:checked={configNoLimit}
class="accent-exo-yellow w-4 h-4"
/>
<span class="text-xs font-mono text-white/80">Unlimited storage</span>
</label>
<!-- Max storage input -->
<div class="space-y-1">
<label
class="text-[11px] font-mono text-white/50 uppercase tracking-wider"
for="storage-max-gb"
>
Max storage (GB)
</label>
<input
id="storage-max-gb"
type="number"
step="1"
min="0"
bind:value={configMaxGb}
disabled={configNoLimit}
class="w-full bg-exo-black/60 border border-exo-medium-gray/30 rounded px-3 py-1.5 text-xs font-mono text-white
focus:outline-none focus:border-exo-yellow/40
disabled:opacity-30 disabled:cursor-not-allowed"
placeholder="e.g. 50"
/>
</div>
<!-- Policy selector -->
<div class="space-y-1.5">
<div
class="text-[11px] font-mono text-white/50 uppercase tracking-wider"
>
Eviction policy
</div>
<div class="flex gap-1">
<button
type="button"
class="flex-1 px-3 py-1.5 rounded text-xs font-mono transition-colors
{configPolicy === 'manual'
? 'bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/40'
: 'bg-exo-black/40 text-white/50 border border-exo-medium-gray/30 hover:text-white/70'}"
onclick={() => (configPolicy = "manual")}
>
Manual
</button>
<button
type="button"
class="flex-1 px-3 py-1.5 rounded text-xs font-mono transition-colors
{configPolicy === 'auto-evict'
? 'bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/40'
: 'bg-exo-black/40 text-white/50 border border-exo-medium-gray/30 hover:text-white/70'}"
onclick={() => (configPolicy = "auto-evict")}
>
Auto-evict
</button>
</div>
<p class="text-[10px] text-white/40 font-mono">
{#if configPolicy === "manual"}
Downloads that exceed the limit are rejected. Delete models
manually.
{:else}
Oldest unused models are automatically removed to make room.
{/if}
</p>
</div>
</div>
<!-- Actions -->
<div class="flex justify-end gap-2 mt-5">
<button
type="button"
class="px-3 py-1.5 rounded text-xs font-mono text-white/50 hover:text-white/70 transition-colors"
onclick={() => (storageConfigNode = null)}
>
Cancel
</button>
<button
type="button"
class="px-3 py-1.5 rounded text-xs font-mono bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/40 hover:bg-exo-yellow/30 transition-colors disabled:opacity-50"
onclick={saveStorageConfig}
disabled={configSaving ||
(!configNoLimit && (configMaxGb == null || configMaxGb <= 0))}
>
{configSaving ? "Saving..." : "Save"}
</button>
</div>
</div>
{/if}
<style>
table {
min-width: max-content;
src/exo/api/main.py
+55
View File
@@ -21,6 +21,7 @@ from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from loguru import logger
from pydantic import Field
from exo.api.adapters.chat_completions import (
chat_request_to_text_generation,
@@ -159,6 +160,7 @@ from exo.shared.types.commands import (
ImageGeneration,
PlaceInstance,
SendInputChunk,
SetStorageConfig,
StartDownload,
TaskCancelled,
TaskFinished,
@@ -173,6 +175,7 @@ from exo.shared.types.events import (
)
from exo.shared.types.memory import Memory
from exo.shared.types.state import State
from exo.shared.types.storage import StorageConfig, StoragePolicy
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -181,8 +184,20 @@ from exo.utils.banner import print_startup_banner
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.disk_event_log import DiskEventLog
from exo.utils.power_sampler import PowerSampler
from exo.utils.pydantic_ext import CamelCaseModel
from exo.utils.task_group import TaskGroup
class SetStorageConfigRequest(CamelCaseModel):
max_storage_gb: Annotated[float, Field(ge=0)] | None = None
storage_policy: StoragePolicy = "manual"
class NodeStorageInfo(CamelCaseModel):
config: StorageConfig
used: Memory
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
@@ -351,6 +366,9 @@ class API:
self.app.get("/events")(self.stream_events)
self.app.post("/download/start")(self.start_download)
self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
self.app.get("/storage")(self.get_storage)
self.app.get("/storage/{node_id}")(self.get_storage_node)
self.app.put("/storage/{node_id}")(self.set_storage_config)
self.app.get("/v1/traces")(self.list_traces)
self.app.post("/v1/traces/delete")(self.delete_traces)
self.app.get("/v1/traces/{task_id}")(self.get_trace)
@@ -1879,6 +1897,43 @@ class API:
await self._send_download(command)
return DeleteDownloadResponse(command_id=command.command_id)
async def get_storage(self) -> dict[str, NodeStorageInfo]:
from exo.shared.storage import calculate_used_storage
result: dict[str, NodeStorageInfo] = {}
for node_id, config in self.state.node_storage_config.items():
downloads = list(self.state.downloads.get(node_id, ()))
used = calculate_used_storage(downloads)
result[node_id] = NodeStorageInfo(config=config, used=used)
return result
async def get_storage_node(self, node_id: NodeId) -> NodeStorageInfo:
from exo.shared.storage import calculate_used_storage
config = self.state.node_storage_config.get(node_id, StorageConfig())
downloads = list(self.state.downloads.get(node_id, ()))
used = calculate_used_storage(downloads)
return NodeStorageInfo(config=config, used=used)
async def set_storage_config(
self, node_id: NodeId, request: SetStorageConfigRequest
) -> dict[str, str]:
max_storage = (
Memory.from_gb(request.max_storage_gb)
if request.max_storage_gb is not None
else None
)
command = SetStorageConfig(
target_node_id=node_id,
max_storage=max_storage,
storage_policy=request.storage_policy,
)
await self.command_sender.send(
ForwarderCommand(origin=self._system_id, command=command)
)
return {"status": "ok", "commandId": str(command.command_id)}
@staticmethod
def _get_trace_path(task_id: str) -> Path:
trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
src/exo/download/coordinator.py
+257 -8
View File
@@ -1,11 +1,18 @@
from __future__ import annotations
import asyncio
import json
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
import aiofiles
import aiofiles.os as aios
import anyio
import tomlkit
from anyio import current_time, to_thread
from loguru import logger
from tomlkit.exceptions import TOMLKitError
from exo.download.download_utils import (
RepoDownloadProgress,
@@ -15,8 +22,19 @@ from exo.download.download_utils import (
resolve_existing_model,
)
from exo.download.shard_downloader import ShardDownloader
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
from exo.shared.constants import (
EXO_CONFIG_FILE,
EXO_DEFAULT_MODELS_DIR,
EXO_MODEL_USAGE_FILE,
EXO_MODELS_READ_ONLY_DIRS,
)
from exo.shared.models.model_cards import ModelId, get_model_cards
from exo.shared.storage import (
calculate_used_storage,
check_storage_quota,
compute_evictions_needed,
get_lru_eviction_candidates,
)
from exo.shared.types.commands import (
CancelDownload,
DeleteDownload,
@@ -26,15 +44,21 @@ from exo.shared.types.commands import (
from exo.shared.types.common import NodeId
from exo.shared.types.events import (
Event,
IndexedEvent,
InstanceCreated,
InstanceDeleted,
NodeDownloadProgress,
StorageConfigUpdated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.storage import StorageConfig
from exo.shared.types.worker.downloads import (
DownloadCompleted,
DownloadFailed,
DownloadOngoing,
DownloadPending,
DownloadProgress,
DownloadRejected,
)
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
from exo.utils.channels import Receiver, Sender
@@ -46,13 +70,20 @@ class DownloadCoordinator:
node_id: NodeId
shard_downloader: ShardDownloader
download_command_receiver: Receiver[ForwarderDownloadCommand]
event_receiver: Receiver[IndexedEvent]
event_sender: Sender[Event]
offline: bool = False
storage_config: StorageConfig = field(default_factory=StorageConfig)
# Local state
download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)
# LRU tracking
_model_last_used: dict[ModelId, datetime] = field(default_factory=dict)
_active_model_ids: set[ModelId] = field(default_factory=set)
_recently_evicted: set[ModelId] = field(default_factory=set)
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
_stopped: anyio.Event = field(init=False, default_factory=anyio.Event)
@@ -127,13 +158,66 @@ class DownloadCoordinator:
logger.info(
f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
)
await self._load_model_usage()
try:
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._emit_existing_download_progress)
tg.start_soon(self._event_watcher)
finally:
self._stopped.set()
async def _event_watcher(self) -> None:
active_instances: dict[str, ModelId] = {}
with self.event_receiver as events:
async for indexed_event in events:
match indexed_event.event:
case StorageConfigUpdated(node_id=node_id) if (
node_id == self.node_id
):
self.storage_config = indexed_event.event.storage_config
await self.clear_rejections()
await self._persist_storage_config(
indexed_event.event.storage_config
)
case InstanceCreated(instance=instance):
active_instances[instance.instance_id] = (
instance.shard_assignments.model_id
)
await self.update_active_models(set(active_instances.values()))
case InstanceDeleted(instance_id=instance_id):
if instance_id in active_instances:
del active_instances[instance_id]
await self.update_active_models(
set(active_instances.values())
)
case _:
pass
@staticmethod
async def _persist_storage_config(config: StorageConfig) -> None:
cfg_path = anyio.Path(EXO_CONFIG_FILE)
await cfg_path.parent.mkdir(parents=True, exist_ok=True)
doc = tomlkit.document()
try:
raw = (await cfg_path.read_bytes()).decode("utf-8")
if raw.strip():
doc = tomlkit.parse(raw)
except (FileNotFoundError, TOMLKitError, UnicodeDecodeError):
pass
if config.max_storage is not None:
doc["max_storage_gb"] = round(config.max_storage.in_gb, 2)
else:
doc.pop("max_storage_gb", None) # pyright: ignore[reportUnknownMemberType]
doc.pop("max_storage_bytes", None) # pyright: ignore[reportUnknownMemberType]
doc["storage_policy"] = config.storage_policy
await cfg_path.write_text(tomlkit.dumps(doc)) # pyright: ignore[reportUnknownMemberType]
logger.debug(f"Persisted storage config to {cfg_path}")
async def shutdown(self) -> None:
self._tg.cancel_tasks()
await self._stopped.wait()
@@ -179,6 +263,22 @@ class DownloadCoordinator:
f"Download for {model_id} already in progress, complete, or failed, skipping"
)
return
# Clear previous DownloadRejected so retries can proceed.
if isinstance(status, DownloadRejected):
if self.storage_config.storage_policy == "auto-evict":
# Emit DownloadPending to update global state — prevents master
# from deleting the instance while auto-eviction is in progress.
pending = DownloadPending(
shard_metadata=status.shard_metadata,
node_id=self.node_id,
model_directory=self._default_model_dir(model_id),
)
self.download_status[model_id] = pending
await self.event_sender.send(
NodeDownloadProgress(download_progress=pending)
)
else:
del self.download_status[model_id]
# Check all model directories for pre-existing complete models
found_path = await to_thread.run_sync(resolve_existing_model, model_id)
@@ -193,6 +293,24 @@ class DownloadCoordinator:
)
return
if self.storage_config.max_storage is not None:
model_size = shard.model_card.storage_size
downloads = list(self.download_status.values())
allowed, reason = check_storage_quota(
model_size, self.storage_config, downloads
)
if not allowed:
if self.storage_config.storage_policy == "auto-evict":
evicted = await self._auto_evict_for_space(shard)
if not evicted:
return
else:
used = calculate_used_storage(downloads)
available_mem = self.storage_config.max_storage - used
await self._reject_download(shard, reason, available_mem)
return
# Emit pending status
progress = DownloadPending(
shard_metadata=shard,
@@ -286,13 +404,15 @@ class DownloadCoordinator:
self._tg.start_soon(download_wrapper, scope)
self.active_downloads[model_id] = scope
async def _delete_download(self, model_id: ModelId) -> None:
async def _delete_download(self, model_id: ModelId) -> bool:
# Protect read-only models from deletion
if model_id in self.download_status:
current = self.download_status[model_id]
if isinstance(current, DownloadCompleted) and current.read_only:
logger.warning(f"Refusing to delete read-only model {model_id}")
return
logger.warning(
f"Refusing to delete read-only model {model_id} (from EXO_MODELS_READ_ONLY_DIRS)"
)
return False
# Cancel if active
if model_id in self.active_downloads:
@@ -303,10 +423,11 @@ class DownloadCoordinator:
logger.info(f"Deleting model files for {model_id}")
deleted = await delete_model(model_id)
if deleted:
logger.info(f"Successfully deleted model {model_id}")
else:
logger.warning(f"Model {model_id} was not found on disk")
if not deleted:
logger.warning(f"Failed to delete model {model_id} from disk")
return False
logger.info(f"Successfully deleted model {model_id}")
# Emit pending status to reset UI state, then remove from local tracking
if model_id in self.download_status:
@@ -321,22 +442,35 @@ class DownloadCoordinator:
)
del self.download_status[model_id]
return True
async def _emit_existing_download_progress(self) -> None:
while True:
try:
logger.debug(
"DownloadCoordinator: Fetching and emitting existing download progress..."
)
# Track which evicted models the scanner still sees on disk
evicted_still_on_disk: set[ModelId] = set()
async for (
_,
progress,
) in self.shard_downloader.get_shard_download_status():
model_id = progress.shard.model_card.model_id
if model_id in self._recently_evicted:
evicted_still_on_disk.add(model_id)
continue
# Active downloads emit progress via the callback — don't overwrite
if model_id in self.active_downloads:
continue
# Don't overwrite rejected status — it should persist until user retries
if isinstance(self.download_status.get(model_id), DownloadRejected):
continue
if progress.status == "complete":
found = await to_thread.run_sync(
resolve_existing_model, model_id
@@ -408,6 +542,8 @@ class DownloadCoordinator:
NodeDownloadProgress(download_progress=path_completed)
)
self._recently_evicted = evicted_still_on_disk
logger.debug(
"DownloadCoordinator: Done emitting existing download progress."
)
@@ -416,3 +552,116 @@ class DownloadCoordinator:
f"DownloadCoordinator: Error emitting existing download progress: {e}"
)
await anyio.sleep(60)
async def _reject_download(
self, shard: ShardMetadata, reason: str, available: Memory
) -> None:
assert self.storage_config.max_storage is not None
model_id = shard.model_card.model_id
rejected = DownloadRejected(
shard_metadata=shard,
node_id=self.node_id,
model_directory=self._default_model_dir(model_id),
reason=reason,
required=shard.model_card.storage_size,
available=available if available.in_bytes > 0 else Memory(),
limit=self.storage_config.max_storage,
)
self.download_status[model_id] = rejected
await self.event_sender.send(NodeDownloadProgress(download_progress=rejected))
async def _auto_evict_for_space(self, shard: ShardMetadata) -> bool:
model_id = shard.model_card.model_id
model_size = shard.model_card.storage_size
assert self.storage_config.max_storage is not None
downloads = list(self.download_status.values())
used = calculate_used_storage(downloads)
raw_available = self.storage_config.max_storage - used
available = raw_available if raw_available.in_bytes >= 0 else Memory()
candidates = get_lru_eviction_candidates(
downloads,
self._model_last_used,
frozenset(self._active_model_ids),
)
to_evict = compute_evictions_needed(model_size, available, candidates)
if to_evict is None:
await self._reject_download(
shard,
"Cannot free enough space even after evicting all eligible models",
available,
)
return False
for evict_model_id in to_evict:
logger.info(
f"Auto-evicting model {evict_model_id} to free space for {model_id}"
)
success = await self._delete_download(evict_model_id)
if not success:
current_used = calculate_used_storage(
list(self.download_status.values())
)
current_available = self.storage_config.max_storage - current_used
await self._reject_download(
shard,
f"Failed to delete model {evict_model_id} from disk",
current_available,
)
return False
self._recently_evicted.add(evict_model_id)
return True
async def clear_rejections(self) -> None:
rejected = [
(model_id, status)
for model_id, status in self.download_status.items()
if isinstance(status, DownloadRejected)
]
for model_id, status in rejected:
logger.info(
f"Clearing DownloadRejected for {model_id} after storage config change"
)
pending = DownloadPending(
shard_metadata=status.shard_metadata,
node_id=self.node_id,
model_directory=self._default_model_dir(model_id),
)
self.download_status[model_id] = pending
await self.event_sender.send(
NodeDownloadProgress(download_progress=pending)
)
async def update_active_models(self, active_model_ids: set[ModelId]) -> None:
new_models = active_model_ids - self._active_model_ids
for mid in new_models:
self._model_last_used[mid] = datetime.now(UTC)
self._active_model_ids = active_model_ids.copy()
if new_models:
await self._persist_model_usage()
async def _persist_model_usage(self) -> None:
try:
await aios.makedirs(EXO_MODEL_USAGE_FILE.parent, exist_ok=True)
data = {mid: ts.isoformat() for mid, ts in self._model_last_used.items()}
async with aiofiles.open(EXO_MODEL_USAGE_FILE, "w") as f:
await f.write(json.dumps(data))
except Exception as e:
logger.warning(f"Failed to persist model usage: {e}")
async def _load_model_usage(self) -> None:
try:
if await aios.path.exists(EXO_MODEL_USAGE_FILE):
async with aiofiles.open(EXO_MODEL_USAGE_FILE, "r") as f:
raw: dict[str, str] = json.loads(await f.read()) # pyright: ignore[reportAny]
self._model_last_used = {
ModelId(k): datetime.fromisoformat(v) for k, v in raw.items()
}
logger.debug(
f"Loaded model usage for {len(self._model_last_used)} models"
)
except Exception as e:
logger.warning(f"Failed to load model usage: {e}")
src/exo/download/tests/test_auto_eviction.py
+409
View File
@@ -0,0 +1,409 @@
"""Tests for auto-eviction in the DownloadCoordinator.
Tests exercise _start_download (the production entry point) to verify that
storage quota checks and LRU eviction work end-to-end through the coordinator.
"""
from datetime import UTC, datetime
from pathlib import Path
from unittest.mock import AsyncMock, patch
from anyio.streams.memory import MemoryObjectStreamState
from exo.download.coordinator import DownloadCoordinator
from exo.download.shard_downloader import NoopShardDownloader
from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.types.commands import ForwarderDownloadCommand
from exo.shared.types.common import NodeId
from exo.shared.types.events import Event, IndexedEvent, NodeDownloadProgress
from exo.shared.types.memory import Memory
from exo.shared.types.storage import StorageConfig
from exo.shared.types.worker.downloads import (
DownloadCompleted,
DownloadPending,
DownloadRejected,
)
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
from exo.utils.channels import Receiver, Sender
MODEL_A = ModelId("org/model-a")
MODEL_B = ModelId("org/model-b")
MODEL_C = ModelId("org/model-c")
MODEL_NEW = ModelId("org/model-new")
NODE_ID = NodeId("test-node")
def _shard(model_id: ModelId, size_gb: float) -> ShardMetadata:
return PipelineShardMetadata(
model_card=ModelCard(
model_id=model_id,
storage_size=Memory.from_gb(size_gb),
n_layers=32,
hidden_size=1000,
supports_tensor=True,
tasks=[ModelTask.TextGeneration],
),
device_rank=0,
world_size=1,
start_layer=0,
end_layer=32,
n_layers=32,
)
def _completed(model_id: ModelId, size_gb: float) -> DownloadCompleted:
return DownloadCompleted(
node_id=NODE_ID,
shard_metadata=_shard(model_id, size_gb),
total=Memory.from_gb(size_gb),
)
def _make_coordinator(
storage_config: StorageConfig,
download_status: dict[ModelId, DownloadCompleted],
model_last_used: dict[ModelId, datetime] | None = None,
) -> tuple[DownloadCoordinator, Receiver[Event]]:
state = MemoryObjectStreamState[Event](max_buffer_size=100)
event_sender = Sender[Event](_state=state)
event_receiver = Receiver[Event](_state=state)
cmd_state: MemoryObjectStreamState[ForwarderDownloadCommand] = (
MemoryObjectStreamState(max_buffer_size=100)
)
cmd_receiver: Receiver[ForwarderDownloadCommand] = Receiver(_state=cmd_state)
idx_state: MemoryObjectStreamState[IndexedEvent] = MemoryObjectStreamState(
max_buffer_size=100
)
idx_receiver: Receiver[IndexedEvent] = Receiver(_state=idx_state)
coordinator = DownloadCoordinator(
node_id=NODE_ID,
shard_downloader=NoopShardDownloader(),
download_command_receiver=cmd_receiver,
event_receiver=idx_receiver,
event_sender=event_sender,
storage_config=storage_config,
)
coordinator.download_status = dict(download_status)
if model_last_used is not None:
coordinator._model_last_used = model_last_used # pyright: ignore[reportPrivateUsage]
return coordinator, event_receiver
async def _start_download(
coordinator: DownloadCoordinator, shard: ShardMetadata
) -> None:
await coordinator._start_download(shard) # pyright: ignore[reportPrivateUsage]
class TestStartDownloadAutoEviction:
"""Tests that go through _start_download — the production entry point."""
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_evicts_oldest_model_to_fit_new_download(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""_start_download should trigger auto-eviction of the oldest model."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
model_last_used={
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
},
)
await _start_download(coordinator, _shard(MODEL_NEW, 5))
# MODEL_A (oldest) should have been evicted
mock_delete.assert_called_once_with(MODEL_A)
assert MODEL_A not in coordinator.download_status
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_evicts_multiple_in_lru_order(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""_start_download evicts multiple models oldest-first until space is freed."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{
MODEL_A: _completed(MODEL_A, 3),
MODEL_B: _completed(MODEL_B, 3),
MODEL_C: _completed(MODEL_C, 3),
},
model_last_used={
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
MODEL_C: datetime(2024, 12, 1, tzinfo=UTC),
},
)
# Need 8 GiB, have 1 GiB free — need to free 7 GiB
await _start_download(coordinator, _shard(MODEL_NEW, 8))
evicted = [call.args[0] for call in mock_delete.call_args_list]
assert evicted == [MODEL_A, MODEL_B, MODEL_C]
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_rejects_when_cannot_free_enough_space(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""_start_download emits DownloadRejected when eviction can't free enough."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 2)},
)
await _start_download(coordinator, _shard(MODEL_NEW, 20))
mock_delete.assert_not_called()
assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected)
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_no_eviction_when_space_available(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""_start_download proceeds without evicting when enough space exists."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 2)},
)
await _start_download(coordinator, _shard(MODEL_NEW, 5))
mock_delete.assert_not_called()
# Download should have started (not rejected)
assert MODEL_NEW in coordinator.download_status
assert not isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected)
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_manual_policy_rejects_instead_of_evicting(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""With manual policy, _start_download rejects instead of auto-evicting."""
config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
)
await _start_download(coordinator, _shard(MODEL_NEW, 5))
mock_delete.assert_not_called()
assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected)
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_eviction_emits_pending_event_for_evicted_model(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""Evicted models emit DownloadPending events (to reset UI state)."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, event_receiver = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
model_last_used={
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
},
)
await _start_download(coordinator, _shard(MODEL_NEW, 5))
events = event_receiver.collect()
pending_events = [
e
for e in events
if isinstance(e, NodeDownloadProgress)
and isinstance(e.download_progress, DownloadPending)
and e.download_progress.shard_metadata.model_card.model_id == MODEL_A
]
assert len(pending_events) == 1
class TestActiveModelProtection:
"""Tests that update_active_models protects models from eviction.
These tests use the public API (update_active_models) to mark models as
active, then trigger eviction through _start_download.
"""
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_active_model_not_evicted(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""A model marked active via update_active_models must not be evicted."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
model_last_used={
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC), # oldest
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
},
)
# Mark MODEL_A as active through the public API
await coordinator.update_active_models({MODEL_A})
await _start_download(coordinator, _shard(MODEL_NEW, 5))
# MODEL_A is active — MODEL_B should be evicted instead
mock_delete.assert_called_once_with(MODEL_B)
assert MODEL_A in coordinator.download_status
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=True,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_all_active_models_rejected(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""When all models are active, eviction is impossible — download is rejected."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
)
await coordinator.update_active_models({MODEL_A, MODEL_B})
await _start_download(coordinator, _shard(MODEL_NEW, 5))
mock_delete.assert_not_called()
assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected)
class TestDiskDeleteFailure:
"""Tests that eviction fails properly when disk delete fails."""
@patch(
"exo.download.coordinator.delete_model",
new_callable=AsyncMock,
return_value=False,
)
@patch("exo.download.coordinator.resolve_model_in_path", return_value=None)
async def test_eviction_rejected_on_disk_delete_failure(
self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
) -> None:
"""When disk delete fails, auto-eviction emits DownloadRejected."""
config = StorageConfig(
max_storage=Memory.from_gb(10), storage_policy="auto-evict"
)
coordinator, _ = _make_coordinator(
config,
{MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
model_last_used={
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
},
)
await _start_download(coordinator, _shard(MODEL_NEW, 5))
# Should have tried to delete oldest model and failed
mock_delete.assert_called_once_with(MODEL_A)
# New model should be rejected
assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected)
# Eviction target should still be in download_status (not removed)
assert MODEL_A in coordinator.download_status
class TestLruPersistence:
"""Tests for _persist_model_usage and _load_model_usage round-trip."""
async def test_persist_then_load_round_trip(self, tmp_path: Path) -> None:
"""Persisting then loading recovers the same data."""
usage_file = tmp_path / "model_usage.json"
coordinator, _ = _make_coordinator(StorageConfig(), {})
coordinator._model_last_used = { # pyright: ignore[reportPrivateUsage]
MODEL_A: datetime(2024, 1, 15, 12, 30, 0, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, 0, 0, 0, tzinfo=UTC),
}
with patch("exo.download.coordinator.EXO_MODEL_USAGE_FILE", usage_file):
await coordinator._persist_model_usage() # pyright: ignore[reportPrivateUsage]
# Create a fresh coordinator and load
coordinator2, _ = _make_coordinator(StorageConfig(), {})
await coordinator2._load_model_usage() # pyright: ignore[reportPrivateUsage]
assert coordinator2._model_last_used == coordinator._model_last_used # pyright: ignore[reportPrivateUsage]
async def test_load_missing_file_returns_empty(self, tmp_path: Path) -> None:
"""Loading when file doesn't exist returns empty dict."""
usage_file = tmp_path / "nonexistent" / "model_usage.json"
coordinator, _ = _make_coordinator(StorageConfig(), {})
with patch("exo.download.coordinator.EXO_MODEL_USAGE_FILE", usage_file):
await coordinator._load_model_usage() # pyright: ignore[reportPrivateUsage]
assert coordinator._model_last_used == {} # pyright: ignore[reportPrivateUsage]
async def test_load_corrupt_json_returns_empty(self, tmp_path: Path) -> None:
"""Loading corrupt JSON logs warning and returns empty dict."""
usage_file = tmp_path / "model_usage.json"
usage_file.write_text("not valid json {{{")
coordinator, _ = _make_coordinator(StorageConfig(), {})
with patch("exo.download.coordinator.EXO_MODEL_USAGE_FILE", usage_file):
await coordinator._load_model_usage() # pyright: ignore[reportPrivateUsage]
assert coordinator._model_last_used == {} # pyright: ignore[reportPrivateUsage]
src/exo/download/tests/test_re_download.py
+3 -1
View File
@@ -19,7 +19,7 @@ from exo.shared.types.commands import (
StartDownload,
)
from exo.shared.types.common import NodeId, SystemId
from exo.shared.types.events import Event, NodeDownloadProgress
from exo.shared.types.events import Event, IndexedEvent, NodeDownloadProgress
from exo.shared.types.memory import Memory
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
@@ -132,6 +132,7 @@ async def test_re_download_after_delete_completes() -> None:
cmd_send: Sender[ForwarderDownloadCommand]
cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
event_send, event_recv = channel[Event]()
_idx_send, idx_recv = channel[IndexedEvent]()
fake_downloader = FakeShardDownloader()
wrapped_downloader = SingletonShardDownloader(fake_downloader)
@@ -139,6 +140,7 @@ async def test_re_download_after_delete_completes() -> None:
node_id=NODE_ID,
shard_downloader=wrapped_downloader,
download_command_receiver=cmd_recv,
event_receiver=idx_recv,
event_sender=event_send,
)
src/exo/main.py
+39 -2
View File
@@ -21,7 +21,10 @@ from exo.shared.constants import EXO_LOG
from exo.shared.election import Election, ElectionResult
from exo.shared.logging import logger_cleanup, logger_setup
from exo.shared.types.common import NodeId, SessionId
from exo.shared.types.memory import Memory
from exo.shared.types.storage import StorageConfig, StoragePolicy
from exo.utils.channels import Receiver, channel
from exo.utils.info_gatherer.info_gatherer import NodeConfig
from exo.utils.pydantic_ext import CamelCaseModel
from exo.utils.task_group import TaskGroup
from exo.worker.main import Worker
@@ -67,14 +70,18 @@ class Node:
logger.info(f"Starting node {node_id}")
storage_config = await cls._load_storage_config(args)
# Create DownloadCoordinator (unless --no-downloads)
if not args.no_downloads:
download_coordinator = DownloadCoordinator(
node_id,
exo_shard_downloader(offline=args.offline),
event_sender=event_router.sender(),
download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
event_receiver=event_router.receiver(),
event_sender=event_router.sender(),
offline=args.offline,
storage_config=storage_config,
)
else:
download_coordinator = None
@@ -165,6 +172,31 @@ class Node:
sys.exit(1)
self._tg.cancel_tasks()
@staticmethod
async def _load_storage_config(args: "Args") -> StorageConfig:
"""Load storage config: start from config.toml, overlay CLI args."""
# Start from config.toml as the base
base_max_storage: Memory | None = None
base_policy: StoragePolicy = "manual"
node_config = await NodeConfig.gather()
if node_config is not None:
if node_config.max_storage_bytes is not None:
base_max_storage = Memory.from_bytes(node_config.max_storage_bytes)
base_policy = node_config.storage_policy
# CLI args override individual fields (non-default values only)
max_storage = (
Memory.from_gb(args.max_storage_gb)
if args.max_storage_gb is not None
else base_max_storage
)
storage_policy = (
args.storage_policy if args.storage_policy != "manual" else base_policy
)
return StorageConfig(max_storage=max_storage, storage_policy=storage_policy)
async def _elect_loop(self):
with self.election_result_receiver as results:
async for result in results:
@@ -229,15 +261,20 @@ class Node:
if result.is_new_master:
if self.download_coordinator:
await self.download_coordinator.shutdown()
storage_config = self.download_coordinator.storage_config
active_model_ids = self.download_coordinator._active_model_ids # pyright: ignore[reportPrivateUsage]
self.download_coordinator = DownloadCoordinator(
self.node_id,
exo_shard_downloader(offline=self.offline),
event_sender=self.event_router.sender(),
download_command_receiver=self.router.receiver(
topics.DOWNLOAD_COMMANDS
),
event_receiver=self.event_router.receiver(),
event_sender=self.event_router.sender(),
offline=self.offline,
storage_config=storage_config,
)
self.download_coordinator._active_model_ids = active_model_ids # pyright: ignore[reportPrivateUsage]
self._tg.start_soon(self.download_coordinator.run)
if self.worker:
await self.worker.shutdown()
src/exo/master/main.py
+54
View File
@@ -24,6 +24,7 @@ from exo.shared.types.commands import (
PlaceInstance,
RequestEventLog,
SendInputChunk,
SetStorageConfig,
TaskCancelled,
TaskFinished,
TestCommand,
@@ -39,8 +40,10 @@ from exo.shared.types.events import (
InputChunkReceived,
InstanceDeleted,
LocalForwarderEvent,
NodeDownloadProgress,
NodeGatheredInfo,
NodeTimedOut,
StorageConfigUpdated,
TaskCreated,
TaskDeleted,
TaskStatusUpdated,
@@ -49,6 +52,7 @@ from exo.shared.types.events import (
TracesMerged,
)
from exo.shared.types.state import State
from exo.shared.types.storage import StorageConfig
from exo.shared.types.tasks import (
ImageEdits as ImageEditsTask,
)
@@ -62,6 +66,7 @@ from exo.shared.types.tasks import (
from exo.shared.types.tasks import (
TextGeneration as TextGenerationTask,
)
from exo.shared.types.worker.downloads import DownloadRejected
from exo.shared.types.worker.instances import InstanceId
from exo.utils.channels import Receiver, Sender
from exo.utils.disk_event_log import DiskEventLog
@@ -96,6 +101,7 @@ class Master:
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
self._expected_ranks: dict[TaskId, set[int]] = {}
self._has_new_rejections = False
async def run(self):
logger.info("Starting Master")
@@ -357,6 +363,16 @@ class Master:
generated_events.append(
CustomModelCardDeleted(model_id=command.model_id)
)
case SetStorageConfig():
generated_events.append(
StorageConfigUpdated(
node_id=command.target_node_id,
storage_config=StorageConfig(
max_storage=command.max_storage,
storage_policy=command.storage_policy,
),
)
)
case RequestEventLog():
# We should just be able to send everything, since other buffers will ignore old messages
# rate limit to 1000 at a time
@@ -384,6 +400,39 @@ class Master:
)
break
# kill instances whose downloads were rejected (storage limit exceeded)
if self._has_new_rejections:
self._has_new_rejections = False
rejected_instances: set[InstanceId] = set()
for instance_id, instance in self.state.instances.items():
for node_id in instance.shard_assignments.node_to_runner:
node_downloads = self.state.downloads.get(node_id, ())
for dp in node_downloads:
if (
isinstance(dp, DownloadRejected)
and dp.shard_metadata.model_card.model_id
== instance.shard_assignments.model_id
):
rejected_instances.add(instance_id)
break
for instance_id in rejected_instances:
logger.info(f"Deleting instance {instance_id} — download rejected")
for task in self.state.tasks.values():
if task.instance_id == instance_id and task.task_status in (
TaskStatus.Pending,
TaskStatus.Running,
):
await self.event_sender.send(
TaskStatusUpdated(
task_id=task.task_id,
task_status=TaskStatus.Failed,
)
)
await self.event_sender.send(
InstanceDeleted(instance_id=instance_id)
)
# time out dead nodes
for node_id, time in self.state.last_seen.items():
now = datetime.now(tz=timezone.utc)
@@ -413,6 +462,11 @@ class Master:
indexed = IndexedEvent(event=event, idx=len(self._event_log))
self.state = apply(self.state, indexed)
if isinstance(event, NodeDownloadProgress) and isinstance(
event.download_progress, DownloadRejected
):
self._has_new_rejections = True
event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage]
if isinstance(event, NodeGatheredInfo):
event.when = str(datetime.now(tz=timezone.utc))
src/exo/shared/apply.py
+31 -1
View File
@@ -18,6 +18,7 @@ from exo.shared.types.events import (
NodeGatheredInfo,
NodeTimedOut,
RunnerStatusUpdated,
StorageConfigUpdated,
TaskAcknowledged,
TaskCreated,
TaskDeleted,
@@ -29,6 +30,7 @@ from exo.shared.types.events import (
TracesCollected,
TracesMerged,
)
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import (
NodeIdentity,
NodeNetworkInfo,
@@ -37,6 +39,7 @@ from exo.shared.types.profiling import (
ThunderboltBridgeStatus,
)
from exo.shared.types.state import State
from exo.shared.types.storage import StorageConfig
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.topology import Connection, RDMAConnection
from exo.shared.types.worker.downloads import DownloadProgress
@@ -95,6 +98,8 @@ def event_apply(event: Event, state: State) -> State:
return apply_topology_edge_created(event, state)
case TopologyEdgeDeleted():
return apply_topology_edge_deleted(event, state)
case StorageConfigUpdated():
return apply_storage_config_updated(event, state)
def apply(state: State, event: IndexedEvent) -> State:
@@ -245,6 +250,11 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
node_rdma_ctl = {
key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
}
node_storage_config = {
key: value
for key, value in state.node_storage_config.items()
if key != event.node_id
}
# Only recompute cycles if the leaving node had TB bridge enabled
leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
leaving_node_had_tb_enabled = (
@@ -267,6 +277,7 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
"node_thunderbolt": node_thunderbolt,
"node_thunderbolt_bridge": node_thunderbolt_bridge,
"node_rdma_ctl": node_rdma_ctl,
"node_storage_config": node_storage_config,
"thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
}
)
@@ -298,7 +309,18 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
case NodeDiskUsage():
update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage}
case NodeConfig():
pass
storage_config = StorageConfig(
max_storage=(
Memory.from_bytes(info.max_storage_bytes)
if info.max_storage_bytes is not None
else None
),
storage_policy=info.storage_policy,
)
update["node_storage_config"] = {
**state.node_storage_config,
event.node_id: storage_config,
}
case MiscData():
current_identity = state.node_identities.get(event.node_id, NodeIdentity())
new_identity = current_identity.model_copy(
@@ -377,6 +399,14 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
return state.model_copy(update=update)
def apply_storage_config_updated(event: StorageConfigUpdated, state: State) -> State:
new_node_storage_config = {
**state.node_storage_config,
event.node_id: event.storage_config,
}
return state.model_copy(update={"node_storage_config": new_node_storage_config})
def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
topology = copy.deepcopy(state.topology)
topology.add_connection(event.conn)
src/exo/shared/constants.py
+2
View File
@@ -94,6 +94,8 @@ EXO_ENABLE_IMAGE_MODELS = (
EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
EXO_MODEL_USAGE_FILE = EXO_DATA_HOME / "model_usage.json"
EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
src/exo/shared/storage.py
+85
View File
@@ -0,0 +1,85 @@
from collections.abc import Mapping, Sequence
from datetime import UTC, datetime
from exo.shared.models.model_cards import ModelId
from exo.shared.types.memory import Memory
from exo.shared.types.storage import StorageConfig
from exo.shared.types.worker.downloads import (
DownloadCompleted,
DownloadOngoing,
DownloadProgress,
)
def calculate_used_storage(downloads: Sequence[DownloadProgress]) -> Memory:
total = Memory()
for dp in downloads:
if isinstance(dp, DownloadCompleted):
total = total + dp.total
elif isinstance(dp, DownloadOngoing):
total = total + dp.download_progress.downloaded
return total
def check_storage_quota(
model_size: Memory,
config: StorageConfig,
downloads: Sequence[DownloadProgress],
) -> tuple[bool, str]:
if config.max_storage is None:
return True, ""
used = calculate_used_storage(downloads)
available = config.max_storage - used
if model_size <= available:
return True, ""
return (
False,
f"Need {model_size.in_gb:.1f} GiB, only {max(0, available.in_gb):.1f} GiB available within {config.max_storage.in_gb:.1f} GiB limit",
)
def get_lru_eviction_candidates(
downloads: Sequence[DownloadProgress],
model_last_used: Mapping[ModelId, datetime],
active_model_ids: frozenset[ModelId],
) -> list[tuple[ModelId, DownloadCompleted]]:
candidates: list[tuple[ModelId, DownloadCompleted]] = []
for dp in downloads:
if not isinstance(dp, DownloadCompleted):
continue
if dp.read_only:
continue
model_id = dp.shard_metadata.model_card.model_id
if model_id in active_model_ids:
continue
candidates.append((model_id, dp))
# Sort by last-used time, oldest first
candidates.sort(
key=lambda item: model_last_used.get(item[0], datetime.min.replace(tzinfo=UTC))
)
return candidates
def compute_evictions_needed(
model_size: Memory,
available: Memory,
candidates: list[tuple[ModelId, DownloadCompleted]],
) -> list[ModelId] | None:
if model_size <= available:
return []
space_needed = model_size - available
freed = Memory()
to_evict: list[ModelId] = []
for model_id, completed in candidates:
to_evict.append(model_id)
freed = freed + completed.total
if freed >= space_needed:
return to_evict
return None
src/exo/shared/tests/test_apply/test_apply_storage.py
+45
View File
@@ -0,0 +1,45 @@
from exo.shared.apply import apply_node_timed_out, apply_storage_config_updated
from exo.shared.types.common import NodeId
from exo.shared.types.events import NodeTimedOut, StorageConfigUpdated
from exo.shared.types.memory import Memory
from exo.shared.types.state import State
from exo.shared.types.storage import StorageConfig
NODE_A = NodeId("node-a")
NODE_B = NodeId("node-b")
def test_storage_config_updated_adds_config() -> None:
state = State()
config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
event = StorageConfigUpdated(node_id=NODE_A, storage_config=config)
new_state = apply_storage_config_updated(event, state)
assert NODE_A in new_state.node_storage_config
assert new_state.node_storage_config[NODE_A].max_storage == Memory.from_gb(10)
assert new_state.node_storage_config[NODE_A].storage_policy == "manual"
def test_storage_config_updated_overwrites_existing() -> None:
config1 = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
state = State(node_storage_config={NODE_A: config1})
config2 = StorageConfig(max_storage=Memory.from_gb(20), storage_policy="auto-evict")
event = StorageConfigUpdated(node_id=NODE_A, storage_config=config2)
new_state = apply_storage_config_updated(event, state)
assert new_state.node_storage_config[NODE_A].max_storage == Memory.from_gb(20)
assert new_state.node_storage_config[NODE_A].storage_policy == "auto-evict"
def test_node_timed_out_cleans_up_storage_config() -> None:
config = StorageConfig(max_storage=Memory.from_gb(10))
state = State(node_storage_config={NODE_A: config, NODE_B: config})
event = NodeTimedOut(node_id=NODE_A)
new_state = apply_node_timed_out(event, state)
assert NODE_A not in new_state.node_storage_config
assert NODE_B in new_state.node_storage_config
src/exo/shared/tests/test_storage.py
+336
View File
@@ -0,0 +1,336 @@
from datetime import UTC, datetime
from exo.shared.models.model_cards import ModelId
from exo.shared.storage import (
calculate_used_storage,
check_storage_quota,
compute_evictions_needed,
get_lru_eviction_candidates,
)
from exo.shared.tests.conftest import get_pipeline_shard_metadata
from exo.shared.types.memory import Memory
from exo.shared.types.storage import StorageConfig
from exo.shared.types.worker.downloads import (
DownloadCompleted,
DownloadOngoing,
DownloadPending,
DownloadProgressData,
)
MODEL_A = ModelId("org/model-a")
MODEL_B = ModelId("org/model-b")
MODEL_C = ModelId("org/model-c")
MODEL_D = ModelId("org/model-d")
NODE_ID = "node-1"
def _completed(
model_id: ModelId, size_gb: float, read_only: bool = False
) -> DownloadCompleted:
shard = get_pipeline_shard_metadata(model_id, device_rank=0)
return DownloadCompleted(
node_id=NODE_ID, # type: ignore[arg-type]
shard_metadata=shard,
total=Memory.from_gb(size_gb),
read_only=read_only,
)
class TestCheckStorageQuota:
def test_unlimited_allows(self) -> None:
config = StorageConfig(max_storage=None)
allowed, _reason = check_storage_quota(Memory.from_gb(10), config, [])
assert allowed is True
assert _reason == ""
def test_under_limit_allows(self) -> None:
config = StorageConfig(max_storage=Memory.from_gb(20))
downloads = [_completed(MODEL_A, 5)]
allowed, _reason = check_storage_quota(Memory.from_gb(10), config, downloads)
assert allowed is True
def test_over_limit_rejects(self) -> None:
config = StorageConfig(max_storage=Memory.from_gb(10))
downloads = [_completed(MODEL_A, 5)]
allowed, reason = check_storage_quota(Memory.from_gb(8), config, downloads)
assert allowed is False
assert "Need" in reason
assert "available" in reason
def test_exact_fit_allows(self) -> None:
config = StorageConfig(max_storage=Memory.from_gb(10))
downloads = [_completed(MODEL_A, 5)]
allowed, _ = check_storage_quota(Memory.from_gb(5), config, downloads)
assert allowed is True
class TestGetLruEvictionCandidates:
def test_excludes_active_models(self) -> None:
downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
last_used = {
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 1, 2, tzinfo=UTC),
}
candidates = get_lru_eviction_candidates(
downloads, last_used, frozenset({MODEL_A})
)
assert len(candidates) == 1
assert candidates[0][0] == MODEL_B
def test_excludes_read_only(self) -> None:
downloads = [_completed(MODEL_A, 5, read_only=True), _completed(MODEL_B, 3)]
candidates = get_lru_eviction_candidates(downloads, {}, frozenset())
assert len(candidates) == 1
assert candidates[0][0] == MODEL_B
def test_sorts_oldest_first(self) -> None:
downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
last_used = {
MODEL_A: datetime(2024, 6, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 1, 1, tzinfo=UTC),
}
candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
assert candidates[0][0] == MODEL_B
assert candidates[1][0] == MODEL_A
def test_models_without_usage_get_min(self) -> None:
downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
last_used = {MODEL_A: datetime(2024, 6, 1, tzinfo=UTC)}
candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
assert candidates[0][0] == MODEL_B # no usage -> datetime.min
class TestComputeEvictionsNeeded:
def test_sufficient_candidates(self) -> None:
candidates = [
(MODEL_A, _completed(MODEL_A, 5)),
(MODEL_B, _completed(MODEL_B, 3)),
]
result = compute_evictions_needed(
Memory.from_gb(6), Memory.from_gb(2), candidates
)
assert result is not None
assert MODEL_A in result
def test_insufficient_candidates(self) -> None:
candidates = [
(MODEL_A, _completed(MODEL_A, 2)),
]
result = compute_evictions_needed(
Memory.from_gb(10), Memory.from_gb(2), candidates
)
assert result is None
def test_no_eviction_needed(self) -> None:
candidates = [(MODEL_A, _completed(MODEL_A, 5))]
result = compute_evictions_needed(
Memory.from_gb(3), Memory.from_gb(5), candidates
)
assert result == []
def test_evicts_in_lru_order(self) -> None:
"""Eviction picks candidates in the order given (oldest first from LRU sort)."""
candidates = [
(MODEL_A, _completed(MODEL_A, 2)), # oldest
(MODEL_B, _completed(MODEL_B, 2)), # newer
(MODEL_C, _completed(MODEL_C, 2)), # newest
]
result = compute_evictions_needed(
Memory.from_gb(3), Memory.from_gb(1), candidates
)
assert result == [MODEL_A]
def test_evicts_multiple_until_enough_space(self) -> None:
"""When one model isn't enough, evicts multiple in LRU order."""
candidates = [
(MODEL_A, _completed(MODEL_A, 1)),
(MODEL_B, _completed(MODEL_B, 1)),
(MODEL_C, _completed(MODEL_C, 1)),
]
# Need 4 GiB, have 1 GiB available — need 3 GiB freed
result = compute_evictions_needed(
Memory.from_gb(4), Memory.from_gb(1), candidates
)
assert result == [MODEL_A, MODEL_B, MODEL_C]
def test_evicts_minimum_needed(self) -> None:
"""Stops evicting as soon as enough space is freed."""
candidates = [
(MODEL_A, _completed(MODEL_A, 3)),
(MODEL_B, _completed(MODEL_B, 3)),
]
# Need 5 GiB, have 2 GiB — need 3 GiB freed. Model A alone suffices.
result = compute_evictions_needed(
Memory.from_gb(5), Memory.from_gb(2), candidates
)
assert result == [MODEL_A]
class TestCalculateUsedStorage:
def test_only_counts_completed_and_ongoing(self) -> None:
"""Pending and rejected downloads should not count toward used storage."""
shard_a = get_pipeline_shard_metadata(MODEL_A, device_rank=0)
shard_b = get_pipeline_shard_metadata(MODEL_B, device_rank=0)
downloads = [
_completed(MODEL_A, 5),
DownloadPending(
node_id=NODE_ID, # type: ignore[arg-type]
shard_metadata=shard_a,
),
DownloadOngoing(
node_id=NODE_ID, # type: ignore[arg-type]
shard_metadata=shard_b,
download_progress=DownloadProgressData(
total=Memory.from_gb(10),
downloaded=Memory.from_gb(3),
downloaded_this_session=Memory.from_gb(3),
completed_files=1,
total_files=5,
speed=0,
eta_ms=0,
files={},
),
),
]
used = calculate_used_storage(downloads)
# 5 GiB completed + 3 GiB ongoing downloaded = 8 GiB
assert abs(used.in_gb - 8.0) < 0.01
def test_empty_downloads(self) -> None:
assert calculate_used_storage([]).in_bytes == 0
class TestGetLruEvictionCandidatesExtended:
def test_excludes_non_completed_downloads(self) -> None:
"""Only DownloadCompleted entries are eviction candidates."""
shard_a = get_pipeline_shard_metadata(MODEL_A, device_rank=0)
downloads = [
_completed(MODEL_B, 3),
DownloadPending(
node_id=NODE_ID, # type: ignore[arg-type]
shard_metadata=shard_a,
),
]
candidates = get_lru_eviction_candidates(downloads, {}, frozenset())
assert len(candidates) == 1
assert candidates[0][0] == MODEL_B
def test_all_active_returns_empty(self) -> None:
"""When all completed models are active, no candidates available."""
downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
candidates = get_lru_eviction_candidates(
downloads, {}, frozenset({MODEL_A, MODEL_B})
)
assert candidates == []
def test_three_models_lru_order(self) -> None:
"""Three models sorted correctly: oldest used first."""
downloads = [
_completed(MODEL_A, 2),
_completed(MODEL_B, 3),
_completed(MODEL_C, 1),
]
last_used = {
MODEL_A: datetime(2024, 3, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_C: datetime(2024, 6, 1, tzinfo=UTC),
}
candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
assert [c[0] for c in candidates] == [MODEL_B, MODEL_A, MODEL_C]
def test_mixed_active_readonly_and_regular(self) -> None:
"""Only non-active, non-read-only completed models are candidates."""
downloads = [
_completed(MODEL_A, 5, read_only=True), # excluded: read-only
_completed(MODEL_B, 3), # excluded: active
_completed(MODEL_C, 2), # candidate
_completed(MODEL_D, 1), # candidate
]
last_used = {
MODEL_C: datetime(2024, 6, 1, tzinfo=UTC),
MODEL_D: datetime(2024, 1, 1, tzinfo=UTC),
}
candidates = get_lru_eviction_candidates(
downloads, last_used, frozenset({MODEL_B})
)
assert [c[0] for c in candidates] == [MODEL_D, MODEL_C]
class TestEndToEndEvictionScenario:
"""Tests that combine LRU candidate selection with eviction computation."""
def test_evicts_oldest_model_to_fit_new_one(self) -> None:
"""10 GiB limit, 3 completed models totaling 9 GiB,
need 3 GiB for new model — should evict the oldest."""
downloads = [
_completed(MODEL_A, 3), # oldest used
_completed(MODEL_B, 3),
_completed(MODEL_C, 3), # newest used
]
last_used = {
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
MODEL_C: datetime(2024, 12, 1, tzinfo=UTC),
}
config = StorageConfig(max_storage=Memory.from_gb(10))
new_model_size = Memory.from_gb(3)
# Step 1: quota check fails
allowed, _ = check_storage_quota(new_model_size, config, downloads)
assert not allowed
# Step 2: get candidates in LRU order
candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
assert candidates[0][0] == MODEL_A # oldest
# Step 3: compute what to evict
used = calculate_used_storage(downloads)
assert config.max_storage is not None
available = config.max_storage - used
to_evict = compute_evictions_needed(new_model_size, available, candidates)
assert to_evict == [MODEL_A]
def test_protects_currently_active_model(self) -> None:
"""Active model should not be evicted even if it's the oldest."""
downloads = [
_completed(MODEL_A, 4), # oldest but active
_completed(MODEL_B, 4), # next oldest, evictable
]
last_used = {
MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
}
config = StorageConfig(max_storage=Memory.from_gb(10))
new_model_size = Memory.from_gb(5)
allowed, _ = check_storage_quota(new_model_size, config, downloads)
assert not allowed
# MODEL_A is active — should be excluded
candidates = get_lru_eviction_candidates(
downloads, last_used, frozenset({MODEL_A})
)
assert len(candidates) == 1
assert candidates[0][0] == MODEL_B
used = calculate_used_storage(downloads)
assert config.max_storage is not None
available = config.max_storage - used
to_evict = compute_evictions_needed(new_model_size, available, candidates)
assert to_evict == [MODEL_B]
def test_cannot_evict_enough_returns_none(self) -> None:
"""When all evictable space isn't enough, returns None."""
downloads = [
_completed(MODEL_A, 4, read_only=True), # can't evict
_completed(MODEL_B, 2), # can evict but only 2 GiB
]
config = StorageConfig(max_storage=Memory.from_gb(10))
new_model_size = Memory.from_gb(8)
candidates = get_lru_eviction_candidates(downloads, {}, frozenset())
used = calculate_used_storage(downloads)
assert config.max_storage is not None
available = config.max_storage - used
to_evict = compute_evictions_needed(new_model_size, available, candidates)
assert to_evict is None
src/exo/utils/info_gatherer/info_gatherer.py
+7
View File
@@ -21,6 +21,7 @@ from exo.shared.types.profiling import (
NetworkInterfaceInfo,
ThunderboltBridgeStatus,
)
from exo.shared.types.storage import StoragePolicy
from exo.shared.types.thunderbolt import (
ThunderboltConnection,
ThunderboltConnectivity,
@@ -294,6 +295,9 @@ class ThunderboltBridgeInfo(TaggedModel):
class NodeConfig(TaggedModel):
"""Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there"""
max_storage_bytes: int | None = None
storage_policy: StoragePolicy = "manual"
@classmethod
async def gather(cls) -> Self | None:
cfg_file = anyio.Path(EXO_CONFIG_FILE)
@@ -303,6 +307,9 @@ class NodeConfig(TaggedModel):
try:
contents = (await f.read()).decode("utf-8")
data = tomllib.loads(contents)
if "max_storage_gb" in data:
gb_value: float = float(data.pop("max_storage_gb")) # pyright: ignore[reportAny]
data["max_storage_bytes"] = round(gb_value * (1024**3))
return cls.model_validate(data)
except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError):
logger.warning("Invalid config file, skipping...")