diff --git a/app/EXO/EXO/Models/ClusterState.swift b/app/EXO/EXO/Models/ClusterState.swift index 1272eeddd..787e1a0aa 100644 --- a/app/EXO/EXO/Models/ClusterState.swift +++ b/app/EXO/EXO/Models/ClusterState.swift @@ -264,7 +264,7 @@ struct NodeDownloadStatus { init?(statusKey: String, payload: NodeDownloadPayload) { guard let nodeId = payload.nodeId else { return nil } self.nodeId = nodeId - self.progress = statusKey == "DownloadOngoing" ? payload.downloadProgress : nil + self.progress = statusKey == "ModelDownloading" ? payload.downloadProgress : nil } } diff --git a/dashboard/src/lib/utils/downloads.ts b/dashboard/src/lib/utils/downloads.ts index 491149109..316a41f9b 100644 --- a/dashboard/src/lib/utils/downloads.ts +++ b/dashboard/src/lib/utils/downloads.ts @@ -5,7 +5,7 @@ * Record> * * Each entry is a tagged union object like: - * { "DownloadCompleted": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } } + * { "ModelReady": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } } */ /** Unwrap one level of tagged-union envelope, returning [tag, payload]. */ @@ -49,7 +49,7 @@ export function extractShardMetadata( return shardMetadata as Record; } -/** Get the download tag (DownloadCompleted, DownloadOngoing, etc.) from a wrapped entry. */ +/** Get the download tag (ModelReady, ModelDownloading, etc.) from a wrapped entry. */ export function getDownloadTag( entry: unknown, ): [string, Record] | null { @@ -73,7 +73,7 @@ function* iterNodeDownloads( } } -/** Check if a specific model is fully downloaded (DownloadCompleted) on a specific node. */ +/** Check if a specific model is fully downloaded (ModelReady) on a specific node. */ export function isModelDownloadedOnNode( downloadsData: Record, nodeId: string, @@ -83,12 +83,12 @@ export function isModelDownloadedOnNode( if (!Array.isArray(nodeDownloads)) return false; for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) { - if (tag === "DownloadCompleted" && entryModelId === modelId) return true; + if (tag === "ModelReady" && entryModelId === modelId) return true; } return false; } -/** Get all node IDs where a model is fully downloaded (DownloadCompleted). */ +/** Get all node IDs where a model is fully downloaded (ModelReady). */ export function getNodesWithModelDownloaded( downloadsData: Record, modelId: string, @@ -122,7 +122,7 @@ export function getShardMetadataForModel( const shard = extractShardMetadata(payload); if (!shard) continue; - if (tag === "DownloadCompleted") return shard; + if (tag === "ModelReady") return shard; if (!fallback) fallback = shard; } } @@ -131,7 +131,7 @@ export function getShardMetadataForModel( /** * Get the download status tag for a specific model on a specific node. - * Returns the "best" status: DownloadCompleted > DownloadOngoing > others. + * Returns the "best" status: ModelReady > ModelDownloading > others. */ export function getModelDownloadStatus( downloadsData: Record, @@ -144,8 +144,8 @@ export function getModelDownloadStatus( let best: string | null = null; for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) { if (entryModelId !== modelId) continue; - if (tag === "DownloadCompleted") return tag; - if (tag === "DownloadOngoing") best = tag; + if (tag === "ModelReady") return tag; + if (tag === "ModelDownloading") best = tag; else if (!best) best = tag; } return best; diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index f54a7d721..79e71310d 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -1582,7 +1582,6 @@ perNode: NodeDownloadStatus[]; failedError: string | null; rejectedError: string | null; - evicted: boolean; } { const empty = { isDownloading: false, @@ -1590,7 +1589,6 @@ perNode: [] as NodeDownloadStatus[], failedError: null, rejectedError: null, - evicted: false, }; if (!downloadsData || Object.keys(downloadsData).length === 0) { @@ -1622,8 +1620,8 @@ const downloadModelId = extractModelIdFromDownload(downloadPayload); if (!downloadModelId || downloadModelId !== modelId) continue; - // DownloadFailed — return with any data collected so far - if (downloadKind === "DownloadFailed") { + // ModelDownloadFailed — return with any data collected so far + if (downloadKind === "ModelDownloadFailed") { return { isDownloading: false, progress: null, @@ -1633,12 +1631,11 @@ (downloadPayload.error_message as string) || "Download failed", rejectedError: null, - evicted: false, }; } - // DownloadRejected — storage limit exceeded - if (downloadKind === "DownloadRejected") { + // ModelRejected — storage limit exceeded + if (downloadKind === "ModelRejected") { return { isDownloading: false, progress: null, @@ -1646,25 +1643,12 @@ failedError: null, rejectedError: (downloadPayload.reason as string) || "Storage limit exceeded", - evicted: false, - }; - } - - // DownloadEvicted — model was evicted from storage - if (downloadKind === "DownloadEvicted") { - return { - isDownloading: false, - progress: null, - perNode: Array.from(perNodeMap.values()), - failedError: null, - rejectedError: null, - evicted: true, }; } if ( - downloadKind !== "DownloadOngoing" && - downloadKind !== "DownloadPending" && + downloadKind !== "ModelDownloading" && + downloadKind !== "ModelNotDownloading" && downloadKind !== "DownloadCompleted" ) continue; @@ -1683,7 +1667,7 @@ continue; } - if (downloadKind === "DownloadPending") { + if (downloadKind === "ModelNotDownloading") { const pendingDownloaded = getBytes( downloadPayload.downloaded ?? downloadPayload.downloaded_bytes ?? @@ -1707,7 +1691,7 @@ continue; } - // DownloadOngoing + // ModelDownloading const progress = parseDownloadProgress(downloadPayload); if ( !progress || @@ -1754,7 +1738,6 @@ perNode, failedError: null, rejectedError: null, - evicted: false, }; } @@ -1776,7 +1759,6 @@ perNode, failedError: null, rejectedError: null, - evicted: false, }; } @@ -1872,17 +1854,6 @@ }; } - if (result.evicted) { - return { - isDownloading: false, - isFailed: false, - errorMessage: null, - progress: null, - statusText: "EVICTED", - perNode: [], - }; - } - if (!result.isDownloading) { const statusInfo = deriveInstanceStatus(instanceWrapped); return { @@ -2584,15 +2555,6 @@ if (prevStatus !== "SHUTDOWN" && currentStatus === "SHUTDOWN") { addToast({ type: "info", message: `Model shut down: ${shortName}` }); } - - // Any -> Evicted - if (prevStatus !== "EVICTED" && currentStatus === "EVICTED") { - addToast({ - type: "info", - message: `Model evicted: ${shortName}`, - duration: 6000, - }); - } } } diff --git a/dashboard/src/routes/downloads/+page.svelte b/dashboard/src/routes/downloads/+page.svelte index f11ed01af..7edc072f4 100644 --- a/dashboard/src/routes/downloads/+page.svelte +++ b/dashboard/src/routes/downloads/+page.svelte @@ -48,7 +48,6 @@ limitBytes: number; modelDirectory?: string; } - | { kind: "evicted"; evictedFor: string; modelDirectory?: string } | { kind: "not_present" }; type ModelCardInfo = { @@ -146,9 +145,9 @@ const tagged = getDownloadTag(entry); if (!tagged) continue; const [tag, payload] = tagged; - if (tag === "DownloadCompleted") { + if (tag === "ModelReady") { total += getBytes(payload.total); - } else if (tag === "DownloadOngoing") { + } else if (tag === "ModelDownloading") { const prog = (payload.download_progress ?? payload.downloadProgress) as | Record | undefined; @@ -169,7 +168,6 @@ downloading: 4, pending: 3, rejected: 2, - evicted: 1, failed: 1, not_present: 0, }; @@ -320,14 +318,14 @@ ((payload.model_directory ?? payload.modelDirectory) as string) || undefined; let cell: CellStatus; - if (tag === "DownloadCompleted") { + if (tag === "ModelReady") { const totalBytes = getBytes(payload.total); cell = { kind: "completed", totalBytes, modelDirectory, }; - } else if (tag === "DownloadOngoing") { + } else if (tag === "ModelDownloading") { const rawProgress = payload.download_progress ?? payload.downloadProgress ?? {}; const prog = rawProgress as Record; @@ -347,7 +345,7 @@ etaMs, modelDirectory, }; - } else if (tag === "DownloadRejected") { + } else if (tag === "ModelRejected") { cell = { kind: "rejected", reason: (payload.reason as string) ?? "Storage limit exceeded", @@ -356,13 +354,8 @@ limitBytes: getBytes(payload.limit), modelDirectory, }; - } else if (tag === "DownloadFailed") { + } else if (tag === "ModelDownloadFailed") { cell = { kind: "failed", modelDirectory }; - } else if (tag === "DownloadEvicted") { - const evictedFor = - ((payload.evicted_for ?? payload.evictedFor) as string) ?? - "unknown"; - cell = { kind: "evicted", evictedFor, modelDirectory }; } else { const downloaded = getBytes( payload.downloaded ?? @@ -822,51 +815,6 @@ {/if} - {:else if cell.kind === "evicted"} -
- - - - - Evicted for {cell.evictedFor.split("/").pop()} - - {#if row.shardMetadata} - - {/if} -
{:else if cell.kind === "failed"}
ModelReady: def _make_coordinator( storage_config: StorageConfig, - download_status: dict[ModelId, ModelReady | ModelEvicted | ModelRejected], + download_status: dict[ModelId, ModelReady | ModelRejected], model_last_used: dict[ModelId, datetime] | None = None, ) -> tuple[DownloadCoordinator, Receiver[Event]]: state = MemoryObjectStreamState[Event](max_buffer_size=100) @@ -129,9 +128,7 @@ class TestStartDownloadAutoEviction: # MODEL_A (oldest) should have been evicted mock_delete.assert_called_once_with(MODEL_A) - evicted_status = coordinator.download_status[MODEL_A] - assert isinstance(evicted_status, ModelEvicted) - assert evicted_status.evicted_for == MODEL_NEW + assert MODEL_A not in coordinator.download_status @patch( "exo.download.coordinator.delete_model", @@ -241,10 +238,10 @@ class TestStartDownloadAutoEviction: return_value=True, ) @patch("exo.download.coordinator.resolve_model_in_path", return_value=None) - async def test_eviction_emits_pending_event_for_evicted_model( + async def test_eviction_emits_not_downloading_event_for_evicted_model( self, _mock_resolve: AsyncMock, mock_delete: AsyncMock ) -> None: - """Evicted models emit DownloadEvicted events directly (no intermediate DownloadPending).""" + """Evicted models emit ModelNotDownloading events and are removed from status.""" config = StorageConfig( max_storage=Memory.from_gb(10), storage_policy="auto-evict" ) @@ -260,17 +257,15 @@ class TestStartDownloadAutoEviction: await _start_download(coordinator, _shard(MODEL_NEW, 5)) events = event_receiver.collect() - evicted_events = [ + eviction_events = [ e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, ModelEvicted) + and isinstance(e.download_progress, ModelNotDownloading) and e.download_progress.shard_metadata.model_card.model_id == MODEL_A ] - assert len(evicted_events) == 1 - evicted_dp = evicted_events[0].download_progress - assert isinstance(evicted_dp, ModelEvicted) - assert evicted_dp.evicted_for == MODEL_NEW + assert len(eviction_events) == 1 + assert MODEL_A not in coordinator.download_status class TestActiveModelProtection: @@ -424,10 +419,10 @@ class TestEvictionEvents: return_value=True, ) @patch("exo.download.coordinator.resolve_model_in_path", return_value=None) - async def test_eviction_emits_download_evicted_event( + async def test_eviction_emits_not_downloading_event( self, _mock_resolve: AsyncMock, mock_delete: AsyncMock ) -> None: - """Eviction emits a DownloadEvicted event with evicted_for set.""" + """Eviction emits a ModelNotDownloading event and removes from status.""" config = StorageConfig( max_storage=Memory.from_gb(10), storage_policy="auto-evict" ) @@ -443,17 +438,15 @@ class TestEvictionEvents: await _start_download(coordinator, _shard(MODEL_NEW, 5)) events = event_receiver.collect() - evicted_events = [ + eviction_events = [ e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, ModelEvicted) + and isinstance(e.download_progress, ModelNotDownloading) + and e.download_progress.shard_metadata.model_card.model_id == MODEL_A ] - assert len(evicted_events) == 1 - evicted_dp = evicted_events[0].download_progress - assert isinstance(evicted_dp, ModelEvicted) - assert evicted_dp.evicted_for == MODEL_NEW - assert evicted_dp.shard_metadata.model_card.model_id == MODEL_A + assert len(eviction_events) == 1 + assert MODEL_A not in coordinator.download_status @patch( "exo.download.coordinator.delete_model", @@ -461,10 +454,10 @@ class TestEvictionEvents: return_value=True, ) @patch("exo.download.coordinator.resolve_model_in_path", return_value=None) - async def test_multi_eviction_emits_evicted_event_per_model( + async def test_multi_eviction_emits_event_per_model( self, _mock_resolve: AsyncMock, _mock_delete: AsyncMock ) -> None: - """Each evicted model gets its own DownloadEvicted event.""" + """Each evicted model gets its own ModelNotDownloading event.""" config = StorageConfig( max_storage=Memory.from_gb(10), storage_policy="auto-evict" ) @@ -485,36 +478,30 @@ class TestEvictionEvents: await _start_download(coordinator, _shard(MODEL_NEW, 8)) events = event_receiver.collect() - evicted_events = [ + eviction_events = [ e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, ModelEvicted) + and isinstance(e.download_progress, ModelNotDownloading) + and e.download_progress.shard_metadata.model_card.model_id != MODEL_NEW ] evicted_model_ids = [ e.download_progress.shard_metadata.model_card.model_id - for e in evicted_events + for e in eviction_events ] assert evicted_model_ids == [MODEL_A, MODEL_B, MODEL_C] - for e in evicted_events: - dp = e.download_progress - assert isinstance(dp, ModelEvicted) - assert dp.evicted_for == MODEL_NEW + for mid in [MODEL_A, MODEL_B, MODEL_C]: + assert mid not in coordinator.download_status class TestClearRejections: - """Tests for clear_rejections behavior with DownloadEvicted.""" + """Tests for clear_rejections behavior.""" - async def test_clear_rejections_preserves_evicted(self) -> None: - """clear_rejections resets DownloadRejected but keeps DownloadEvicted.""" + async def test_clear_rejections_resets_rejected(self) -> None: + """clear_rejections resets ModelRejected to ModelNotDownloading.""" config = StorageConfig( max_storage=Memory.from_gb(10), storage_policy="auto-evict" ) - evicted = ModelEvicted( - node_id=NODE_ID, - shard_metadata=_shard(MODEL_A, 4), - evicted_for=MODEL_NEW, - ) rejected = ModelRejected( node_id=NODE_ID, shard_metadata=_shard(MODEL_B, 4), @@ -525,14 +512,14 @@ class TestClearRejections: ) coordinator, _ = _make_coordinator( config, - {MODEL_A: evicted, MODEL_B: rejected}, + {MODEL_A: _completed(MODEL_A, 4), MODEL_B: rejected}, ) await coordinator.clear_rejections() - # Evicted should remain unchanged - assert isinstance(coordinator.download_status[MODEL_A], ModelEvicted) - # Rejected should be cleared to Pending + # Completed should remain unchanged + assert isinstance(coordinator.download_status[MODEL_A], ModelReady) + # Rejected should be cleared assert isinstance(coordinator.download_status[MODEL_B], ModelNotDownloading) async def test_clear_rejections_on_policy_only_change(self) -> None: diff --git a/src/exo/shared/storage.py b/src/exo/shared/storage.py index 28214bffd..a592734d0 100644 --- a/src/exo/shared/storage.py +++ b/src/exo/shared/storage.py @@ -22,8 +22,8 @@ from exo.shared.types.storage import ( ) from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import ( - ModelReady, ModelDownloading, + ModelReady, ModelStatus, ) from exo.shared.types.worker.instances import Instance, InstanceId diff --git a/src/exo/shared/tests/test_storage.py b/src/exo/shared/tests/test_storage.py index 36409bb40..b4f50eb51 100644 --- a/src/exo/shared/tests/test_storage.py +++ b/src/exo/shared/tests/test_storage.py @@ -25,10 +25,10 @@ from exo.shared.types.storage import ( ) from exo.shared.types.tasks import LoadModel, TaskId, TaskStatus from exo.shared.types.worker.downloads import ( - ModelReady, + DownloadProgressData, ModelDownloading, ModelNotDownloading, - DownloadProgressData, + ModelReady, ) from exo.shared.types.worker.instances import InstanceId, MlxRingInstance from exo.shared.types.worker.runners import RunnerId, ShardAssignments diff --git a/src/exo/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py index 524620867..24e367f56 100644 --- a/src/exo/shared/types/worker/downloads.py +++ b/src/exo/shared/types/worker/downloads.py @@ -3,7 +3,7 @@ from typing import Literal from pydantic import BaseModel, ConfigDict, Field, PositiveInt -from exo.shared.types.common import ModelId, NodeId +from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.worker.shards import ShardMetadata from exo.utils.pydantic_ext import FrozenModel, TaggedModel @@ -54,17 +54,12 @@ class ModelRejected(BaseModelStatus): limit: Memory -class ModelEvicted(BaseModelStatus): - evicted_for: ModelId - - ModelStatus = ( ModelNotDownloading | ModelReady | ModelDownloadFailed | ModelDownloading | ModelRejected - | ModelEvicted ) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 2d5fc5e88..3a44285bd 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -21,9 +21,9 @@ from exo.shared.types.tasks import ( ) from exo.shared.types.text_generation import Base64Image, Base64ImageHash from exo.shared.types.worker.downloads import ( - ModelReady, ModelDownloadFailed, ModelDownloading, + ModelReady, ModelStatus, ) from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId diff --git a/tmp/old_tests/get_all_models_on_cluster.py b/tmp/old_tests/get_all_models_on_cluster.py index d150e9ea3..5455b7834 100755 --- a/tmp/old_tests/get_all_models_on_cluster.py +++ b/tmp/old_tests/get_all_models_on_cluster.py @@ -19,7 +19,7 @@ with urlopen(f"http://{ip}:52415/state", timeout=5) as r: def mid(x: dict[str, Any]) -> str | None: for k in ( - "DownloadCompleted", + "ModelReady", "shardMetadata", "PipelineShardMetadata", "modelCard", diff --git a/tools/src/exo_tools/harness.py b/tools/src/exo_tools/harness.py index a80da6254..b61aecfd4 100644 --- a/tools/src/exo_tools/harness.py +++ b/tools/src/exo_tools/harness.py @@ -331,8 +331,8 @@ def run_planning_phase( node_downloads = client.get_node_downloads(node_id) or [] already_downloaded = any( - "DownloadCompleted" in p - and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ + "ModelReady" in p + and unwrap_instance(p["ModelReady"]["shardMetadata"])["modelCard"][ "modelId" ] == full_model_id @@ -370,14 +370,13 @@ def run_planning_phase( completed = [ ( - unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ + unwrap_instance(p["ModelReady"]["shardMetadata"])["modelCard"][ "modelId" ], - p["DownloadCompleted"]["total"]["inBytes"], + p["ModelReady"]["total"]["inBytes"], ) for p in node_downloads - if "DownloadCompleted" in p - and not p["DownloadCompleted"].get("readOnly", False) + if "ModelReady" in p and not p["ModelReady"].get("readOnly", False) ] for del_model, size in sorted(completed, key=lambda x: x[1]): logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)") @@ -410,20 +409,20 @@ def run_planning_phase( for node_id in node_ids: node_downloads = client.get_node_downloads(node_id) or [] done = any( - "DownloadCompleted" in p - and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[ - "modelCard" - ]["modelId"] + "ModelReady" in p + and unwrap_instance(p["ModelReady"]["shardMetadata"])["modelCard"][ + "modelId" + ] == full_model_id for p in node_downloads ) failed = [ - p["DownloadFailed"]["errorMessage"] + p["ModelDownloadFailed"]["errorMessage"] for p in node_downloads - if "DownloadFailed" in p - and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][ - "modelId" - ] + if "ModelDownloadFailed" in p + and unwrap_instance(p["ModelDownloadFailed"]["shardMetadata"])[ + "modelCard" + ]["modelId"] == full_model_id ] if failed: