Remove ModelEvicted status

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
     }
 }
 

dashboard/src/lib/utils/downloads.ts

@@ -5,7 +5,7 @@
  *   Record<NodeId, Array<TaggedDownloadEntry>>
  *
  * 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<string, unknown>;
 }
 
-/** 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<string, unknown>] | 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<string, unknown[]>,
   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<string, unknown[]>,
   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<string, unknown[]>,
@@ -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;

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,
-          });
-        }
       }
     }
 

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<string, unknown>
           | 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<string, unknown>;
@@ -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 @@
                           </button>
                         {/if}
                       </div>
-                    {:else if cell.kind === "evicted"}
-                      <div
-                        class="flex flex-col items-center gap-1"
-                        title="Evicted for {cell.evictedFor}"
-                      >
-                        <svg
-                          class="w-5 h-5 text-blue-400"
-                          viewBox="0 0 20 20"
-                          fill="currentColor"
-                        >
-                          <path
-                            fill-rule="evenodd"
-                            d="M18 10a8 8 0 11-16 0 8 8 0 0116 0zm-7-4a1 1 0 11-2 0 1 1 0 012 0zM9 9a1 1 0 000 2v3a1 1 0 001 1h1a1 1 0 100-2v-3a1 1 0 00-1-1H9z"
-                            clip-rule="evenodd"
-                          ></path>
-                        </svg>
-                        <span
-                          class="text-[10px] text-blue-400/80 leading-tight text-center max-w-[100px] truncate"
-                        >
-                          Evicted for {cell.evictedFor.split("/").pop()}
-                        </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="Re-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"

src/exo/download/coordinator.py

@@ -57,7 +57,6 @@ from exo.shared.types.storage import (
 from exo.shared.types.worker.downloads import (
     ModelDownloadFailed,
     ModelDownloading,
-    ModelEvicted,
     ModelNotDownloading,
     ModelReady,
     ModelRejected,
@@ -81,6 +80,7 @@ class DownloadCoordinator:
     # Local state
     download_status: dict[ModelId, ModelStatus] = field(default_factory=dict)
     active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)
+    _deleting: set[ModelId] = field(default_factory=set)
 
     _model_last_used: dict[ModelId, datetime] = field(default_factory=dict)
     _active_model_ids: set[ModelId] = field(default_factory=set)
@@ -441,9 +441,8 @@ class DownloadCoordinator:
                 ) in self.shard_downloader.get_shard_download_status():
                     model_id = progress.shard.model_card.model_id
 
-                    # Don't overwrite ModelEvicted — the model may still be on disk
-                    # while deletion is finishing
-                    if isinstance(self.download_status.get(model_id), ModelEvicted):
+                    # Don't overwrite status while deletion is in progress
+                    if model_id in self._deleting:
                         continue
 
                     # Active downloads emit progress via the callback — don't overwrite
@@ -588,7 +587,12 @@ class DownloadCoordinator:
                 f"Auto-evicting model {evict_model_id} to free space for {target_model_id}"
             )
             evicted_status = self.download_status.get(evict_model_id)
-            success = await self._remove_model_from_disk(evict_model_id)
+            self._deleting.add(evict_model_id)
+            try:
+                success = await self._remove_model_from_disk(evict_model_id)
+            finally:
+                self._deleting.discard(evict_model_id)
+
             if not success:
                 current_used = calculate_used_storage(
                     list(self.download_status.values())
@@ -603,16 +607,15 @@ class DownloadCoordinator:
                 return False
 
             if evicted_status is not None:
-                evicted = ModelEvicted(
+                not_downloading = ModelNotDownloading(
                     shard_metadata=evicted_status.shard_metadata,
                     node_id=self.node_id,
                     model_directory=self._default_model_dir(evict_model_id),
-                    evicted_for=target_model_id,
                 )
-                self.download_status[evict_model_id] = evicted
                 await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=evicted)
+                    NodeDownloadProgress(download_progress=not_downloading)
                 )
+                del self.download_status[evict_model_id]
 
         return True
 

src/exo/download/tests/test_auto_eviction.py

@@ -19,9 +19,8 @@ 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 (
-    ModelReady,
-    ModelEvicted,
     ModelNotDownloading,
+    ModelReady,
     ModelRejected,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
@@ -62,7 +61,7 @@ def _completed(model_id: ModelId, size_gb: float) -> 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:

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

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

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
 )
 
 

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

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",

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: