diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index a5508f92..8550d5dd 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -261,20 +261,7 @@ class DownloadCoordinator: 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] + 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) @@ -400,7 +387,7 @@ class DownloadCoordinator: self._tg.start_soon(download_wrapper, scope) self.active_downloads[model_id] = scope - async def _delete_download(self, model_id: ModelId) -> bool: + async def _remove_model_from_disk(self, model_id: ModelId) -> bool: # Protect read-only models from deletion if model_id in self.download_status: current = self.download_status[model_id] @@ -424,6 +411,12 @@ class DownloadCoordinator: return False logger.info(f"Successfully deleted model {model_id}") + return True + + async def _delete_download(self, model_id: ModelId) -> bool: + success = await self._remove_model_from_disk(model_id) + if not success: + return False # Emit pending status to reset UI state, then remove from local tracking if model_id in self.download_status: @@ -592,9 +585,8 @@ class DownloadCoordinator: logger.info( f"Auto-evicting model {evict_model_id} to free space for {model_id}" ) - # Capture shard_metadata before _delete_download removes it evicted_status = self.download_status.get(evict_model_id) - success = await self._delete_download(evict_model_id) + success = await self._remove_model_from_disk(evict_model_id) if not success: current_used = calculate_used_storage( list(self.download_status.values()) @@ -607,13 +599,11 @@ class DownloadCoordinator: ) return False - # Overwrite the DownloadPending that _delete_download emitted - # with an explicit DownloadEvicted status if evicted_status is not None: evicted = DownloadEvicted( shard_metadata=evicted_status.shard_metadata, node_id=self.node_id, - model_directory=self._model_dir(evict_model_id), + model_directory=self._default_model_dir(evict_model_id), evicted_for=model_id, ) self.download_status[evict_model_id] = evicted diff --git a/src/exo/download/tests/test_auto_eviction.py b/src/exo/download/tests/test_auto_eviction.py index 82032d86..6f6a0551 100644 --- a/src/exo/download/tests/test_auto_eviction.py +++ b/src/exo/download/tests/test_auto_eviction.py @@ -62,7 +62,9 @@ def _completed(model_id: ModelId, size_gb: float) -> DownloadCompleted: def _make_coordinator( storage_config: StorageConfig, - download_status: dict[ModelId, DownloadCompleted], + download_status: dict[ + ModelId, DownloadCompleted | DownloadEvicted | DownloadRejected + ], model_last_used: dict[ModelId, datetime] | None = None, ) -> tuple[DownloadCoordinator, Receiver[Event]]: state = MemoryObjectStreamState[Event](max_buffer_size=100) @@ -244,7 +246,7 @@ class TestStartDownloadAutoEviction: 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).""" + """Evicted models emit DownloadEvicted events directly (no intermediate DownloadPending).""" config = StorageConfig( max_storage=Memory.from_gb(10), storage_policy="auto-evict" ) @@ -260,14 +262,17 @@ class TestStartDownloadAutoEviction: await _start_download(coordinator, _shard(MODEL_NEW, 5)) events = event_receiver.collect() - pending_events = [ + evicted_events = [ e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, DownloadPending) + and isinstance(e.download_progress, DownloadEvicted) and e.download_progress.shard_metadata.model_card.model_id == MODEL_A ] - assert len(pending_events) == 1 + assert len(evicted_events) == 1 + evicted_dp = evicted_events[0].download_progress + assert isinstance(evicted_dp, DownloadEvicted) + assert evicted_dp.evicted_for == MODEL_NEW class TestActiveModelProtection: @@ -410,3 +415,124 @@ class TestLruPersistence: await coordinator._load_model_usage() # pyright: ignore[reportPrivateUsage] assert coordinator._model_last_used == {} # pyright: ignore[reportPrivateUsage] + + +class TestEvictionEvents: + """Tests that eviction emits the correct sequence of events.""" + + @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_download_evicted_event( + self, _mock_resolve: AsyncMock, mock_delete: AsyncMock + ) -> None: + """Eviction emits a DownloadEvicted event with evicted_for set.""" + 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() + evicted_events = [ + e + for e in events + if isinstance(e, NodeDownloadProgress) + and isinstance(e.download_progress, DownloadEvicted) + ] + assert len(evicted_events) == 1 + evicted_dp = evicted_events[0].download_progress + assert isinstance(evicted_dp, DownloadEvicted) + assert evicted_dp.evicted_for == MODEL_NEW + assert evicted_dp.shard_metadata.model_card.model_id == MODEL_A + + @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_multi_eviction_emits_evicted_event_per_model( + self, _mock_resolve: AsyncMock, _mock_delete: AsyncMock + ) -> None: + """Each evicted model gets its own DownloadEvicted event.""" + config = StorageConfig( + max_storage=Memory.from_gb(10), storage_policy="auto-evict" + ) + coordinator, event_receiver = _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), + }, + ) + + await _start_download(coordinator, _shard(MODEL_NEW, 8)) + + events = event_receiver.collect() + evicted_events = [ + e + for e in events + if isinstance(e, NodeDownloadProgress) + and isinstance(e.download_progress, DownloadEvicted) + ] + evicted_model_ids = [ + e.download_progress.shard_metadata.model_card.model_id + for e in evicted_events + ] + assert evicted_model_ids == [MODEL_A, MODEL_B, MODEL_C] + for e in evicted_events: + dp = e.download_progress + assert isinstance(dp, DownloadEvicted) + assert dp.evicted_for == MODEL_NEW + + +class TestClearRejections: + """Tests for clear_rejections behavior with DownloadEvicted.""" + + async def test_clear_rejections_preserves_evicted(self) -> None: + """clear_rejections resets DownloadRejected but keeps DownloadEvicted.""" + config = StorageConfig( + max_storage=Memory.from_gb(10), storage_policy="auto-evict" + ) + evicted = DownloadEvicted( + node_id=NODE_ID, + shard_metadata=_shard(MODEL_A, 4), + evicted_for=MODEL_NEW, + ) + rejected = DownloadRejected( + node_id=NODE_ID, + shard_metadata=_shard(MODEL_B, 4), + reason="Not enough space", + required=Memory.from_gb(4), + available=Memory.from_gb(1), + limit=Memory.from_gb(10), + ) + coordinator, _ = _make_coordinator( + config, + {MODEL_A: evicted, MODEL_B: rejected}, + ) + + await coordinator.clear_rejections() + + # Evicted should remain unchanged + assert isinstance(coordinator.download_status[MODEL_A], DownloadEvicted) + # Rejected should be cleared to Pending + assert isinstance(coordinator.download_status[MODEL_B], DownloadPending) diff --git a/src/exo/main.py b/src/exo/main.py index 33e0ea0e..6e6def4f 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -176,7 +176,9 @@ class Node: async def _load_storage_config(args: "Args") -> StorageConfig: """Load storage config: start from config.toml, overlay CLI args.""" node_config = await NodeConfig.gather() - base = node_config.storage_config if node_config is not None else StorageConfig() + base = ( + node_config.storage_config if node_config is not None else StorageConfig() + ) # CLI args override individual fields (non-default values only) max_storage = ( @@ -185,7 +187,9 @@ class Node: else base.max_storage ) storage_policy = ( - args.storage_policy if args.storage_policy != "manual" else base.storage_policy + args.storage_policy + if args.storage_policy is not None + else base.storage_policy ) return StorageConfig(max_storage=max_storage, storage_policy=storage_policy) @@ -348,7 +352,7 @@ class Args(CamelCaseModel): bootstrap_peers: list[str] = [] libp2p_port: int max_storage_gb: float | None = None - storage_policy: StoragePolicy = "manual" + storage_policy: StoragePolicy | None = None @classmethod def parse(cls) -> Self: @@ -447,8 +451,8 @@ class Args(CamelCaseModel): "--storage-policy", choices=["manual", "auto-evict"], dest="storage_policy", - default="manual", - help="Storage policy: 'manual' rejects on exceed, 'auto-evict' removes LRU models", + default=None, + help="Storage policy: 'manual' rejects on exceed, 'auto-evict' removes LRU models (default: manual)", ) args = parser.parse_args() diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py index 4a594e02..315cc943 100644 --- a/src/exo/utils/info_gatherer/info_gatherer.py +++ b/src/exo/utils/info_gatherer/info_gatherer.py @@ -15,13 +15,13 @@ from pydantic import ValidationError from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR from exo.shared.types.memory import Memory -from exo.shared.types.storage import StorageConfig from exo.shared.types.profiling import ( DiskUsage, MemoryUsage, NetworkInterfaceInfo, ThunderboltBridgeStatus, ) +from exo.shared.types.storage import StorageConfig from exo.shared.types.thunderbolt import ( ThunderboltConnection, ThunderboltConnectivity,