diff --git a/src/exo/api/main.py b/src/exo/api/main.py index a3972bd4..7693681d 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -178,7 +178,7 @@ 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.downloads import ModelReady from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding from exo.utils.banner import print_startup_banner @@ -1573,7 +1573,7 @@ class API: downloaded_model_ids: set[str] = set() for node_downloads in self.state.downloads.values(): for dl in node_downloads: - if isinstance(dl, DownloadCompleted): + if isinstance(dl, ModelReady): downloaded_model_ids.add(dl.shard_metadata.model_card.model_id) cards = [ @@ -1660,7 +1660,7 @@ class API: downloaded_model_ids: set[str] = set() for node_downloads in self.state.downloads.values(): for dl in node_downloads: - if isinstance(dl, DownloadCompleted): + if isinstance(dl, ModelReady): downloaded_model_ids.add(dl.shard_metadata.model_card.model_id) cards = [c for c in cards if c.model_id in downloaded_model_ids] diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index a942b04a..11852393 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -54,13 +54,13 @@ from exo.shared.types.storage import ( StorageReject, ) from exo.shared.types.worker.downloads import ( - DownloadCompleted, - DownloadEvicted, - DownloadFailed, - DownloadOngoing, - DownloadPending, - DownloadProgress, - DownloadRejected, + ModelDownloadFailed, + ModelDownloading, + ModelEvicted, + ModelNotDownloading, + ModelReady, + ModelRejected, + ModelStatus, ) from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata from exo.utils.channels import Receiver, Sender @@ -78,7 +78,7 @@ class DownloadCoordinator: storage_config: StorageConfig = field(default_factory=StorageConfig) # Local state - download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict) + download_status: dict[ModelId, ModelStatus] = field(default_factory=dict) active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict) _model_last_used: dict[ModelId, datetime] = field(default_factory=dict) @@ -140,7 +140,7 @@ class DownloadCoordinator: and current_time() - self._last_progress_time.get(model_id, 0.0) > throttle_interval_secs ): - ongoing = DownloadOngoing( + ongoing = ModelDownloading( node_id=self.node_id, shard_metadata=callback_shard, download_progress=map_repo_download_progress_to_download_progress_data( @@ -216,7 +216,7 @@ class DownloadCoordinator: logger.info(f"Cancelling download for {model_id}") self.active_downloads[model_id].cancel() current_status = self.download_status[model_id] - pending = DownloadPending( + pending = ModelNotDownloading( shard_metadata=current_status.shard_metadata, node_id=self.node_id, model_directory=self._default_model_dir(model_id), @@ -232,12 +232,12 @@ class DownloadCoordinator: # Check if already downloading, complete, or recently failed if model_id in self.download_status: status = self.download_status[model_id] - if isinstance(status, (DownloadOngoing, DownloadCompleted, DownloadFailed)): + if isinstance(status, (ModelDownloading, ModelReady, ModelDownloadFailed)): logger.debug( f"Download for {model_id} already in progress, complete, or failed, skipping" ) return - if isinstance(status, DownloadRejected): + if isinstance(status, ModelRejected): del self.download_status[model_id] # Check all model directories for pre-existing complete models @@ -271,7 +271,7 @@ class DownloadCoordinator: pass # Emit pending status - progress = DownloadPending( + progress = ModelNotDownloading( shard_metadata=shard, node_id=self.node_id, model_directory=self._default_model_dir(model_id), @@ -307,7 +307,7 @@ class DownloadCoordinator: logger.warning( f"Offline mode: model {model_id} is not fully available locally, cannot download" ) - failed = DownloadFailed( + failed = ModelDownloadFailed( shard_metadata=shard, node_id=self.node_id, error_message=f"Model files not found locally in offline mode: {model_id}", @@ -326,7 +326,7 @@ class DownloadCoordinator: model_id = shard.model_card.model_id # Emit ongoing status - status = DownloadOngoing( + status = ModelDownloading( node_id=self.node_id, shard_metadata=shard, download_progress=map_repo_download_progress_to_download_progress_data( @@ -343,7 +343,7 @@ class DownloadCoordinator: await self.shard_downloader.ensure_shard(shard) except Exception as e: logger.error(f"Download failed for {model_id}: {e}") - failed = DownloadFailed( + failed = ModelDownloadFailed( shard_metadata=shard, node_id=self.node_id, error_message=str(e), @@ -367,7 +367,7 @@ class DownloadCoordinator: # 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: + if isinstance(current, ModelReady) and current.read_only: logger.warning( f"Refusing to delete read-only model {model_id} (from EXO_MODELS_READ_ONLY_DIRS)" ) @@ -397,7 +397,7 @@ class DownloadCoordinator: # Emit pending status to reset UI state, then remove from local tracking if model_id in self.download_status: current_status = self.download_status[model_id] - pending = DownloadPending( + pending = ModelNotDownloading( shard_metadata=current_status.shard_metadata, node_id=self.node_id, model_directory=self._default_model_dir(model_id), @@ -424,14 +424,14 @@ class DownloadCoordinator: # Don't overwrite DownloadEvicted — the model may still be on disk # while deletion is finishing - if isinstance(self.download_status.get(model_id), DownloadEvicted): + if isinstance(self.download_status.get(model_id), ModelEvicted): continue # Active downloads emit progress via the callback — don't overwrite if model_id in self.active_downloads: continue - if isinstance(self.download_status.get(model_id), DownloadRejected): + if isinstance(self.download_status.get(model_id), ModelRejected): continue if progress.status == "complete": @@ -451,7 +451,7 @@ class DownloadCoordinator: ) elif progress.status in ["in_progress", "not_started"]: if progress.downloaded_this_session.in_bytes == 0: - status = DownloadPending( + status = ModelNotDownloading( node_id=self.node_id, shard_metadata=progress.shard, model_directory=self._default_model_dir(model_id), @@ -459,7 +459,7 @@ class DownloadCoordinator: total=progress.total, ) else: - status = DownloadOngoing( + status = ModelDownloading( node_id=self.node_id, shard_metadata=progress.shard, download_progress=map_repo_download_progress_to_download_progress_data( @@ -482,7 +482,7 @@ class DownloadCoordinator: continue if isinstance( self.download_status.get(mid), - (DownloadCompleted, DownloadOngoing, DownloadFailed), + (ModelReady, ModelDownloading, ModelDownloadFailed), ): continue found = await to_thread.run_sync(resolve_existing_model, mid) @@ -519,7 +519,7 @@ class DownloadCoordinator: ) -> None: assert self.storage_config.max_storage is not None model_id = shard.model_card.model_id - rejected = DownloadRejected( + rejected = ModelRejected( shard_metadata=shard, node_id=self.node_id, model_directory=self._default_model_dir(model_id), @@ -557,7 +557,7 @@ class DownloadCoordinator: return False if evicted_status is not None: - evicted = DownloadEvicted( + evicted = ModelEvicted( shard_metadata=evicted_status.shard_metadata, node_id=self.node_id, model_directory=self._default_model_dir(evict_model_id), @@ -574,13 +574,13 @@ class DownloadCoordinator: rejected = [ (model_id, status) for model_id, status in self.download_status.items() - if isinstance(status, DownloadRejected) + if isinstance(status, ModelRejected) ] for model_id, status in rejected: logger.info( f"Clearing DownloadRejected for {model_id} after storage config change" ) - pending = DownloadPending( + pending = ModelNotDownloading( shard_metadata=status.shard_metadata, node_id=self.node_id, model_directory=self._default_model_dir(model_id), diff --git a/src/exo/download/tests/test_auto_eviction.py b/src/exo/download/tests/test_auto_eviction.py index 2e95867a..68b57fc6 100644 --- a/src/exo/download/tests/test_auto_eviction.py +++ b/src/exo/download/tests/test_auto_eviction.py @@ -19,10 +19,10 @@ 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, - DownloadEvicted, - DownloadPending, - DownloadRejected, + ModelReady, + ModelEvicted, + ModelNotDownloading, + ModelRejected, ) from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata from exo.utils.channels import Receiver, Sender @@ -52,8 +52,8 @@ def _shard(model_id: ModelId, size_gb: float) -> ShardMetadata: ) -def _completed(model_id: ModelId, size_gb: float) -> DownloadCompleted: - return DownloadCompleted( +def _completed(model_id: ModelId, size_gb: float) -> ModelReady: + return ModelReady( node_id=NODE_ID, shard_metadata=_shard(model_id, size_gb), total=Memory.from_gb(size_gb), @@ -62,9 +62,7 @@ def _completed(model_id: ModelId, size_gb: float) -> DownloadCompleted: def _make_coordinator( storage_config: StorageConfig, - download_status: dict[ - ModelId, DownloadCompleted | DownloadEvicted | DownloadRejected - ], + download_status: dict[ModelId, ModelReady | ModelEvicted | ModelRejected], model_last_used: dict[ModelId, datetime] | None = None, ) -> tuple[DownloadCoordinator, Receiver[Event]]: state = MemoryObjectStreamState[Event](max_buffer_size=100) @@ -132,7 +130,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, DownloadEvicted) + assert isinstance(evicted_status, ModelEvicted) assert evicted_status.evicted_for == MODEL_NEW @patch( @@ -189,7 +187,7 @@ class TestStartDownloadAutoEviction: await _start_download(coordinator, _shard(MODEL_NEW, 20)) mock_delete.assert_not_called() - assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected) + assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected) @patch( "exo.download.coordinator.delete_model", @@ -214,7 +212,7 @@ class TestStartDownloadAutoEviction: 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) + assert not isinstance(coordinator.download_status[MODEL_NEW], ModelRejected) @patch( "exo.download.coordinator.delete_model", @@ -235,7 +233,7 @@ class TestStartDownloadAutoEviction: await _start_download(coordinator, _shard(MODEL_NEW, 5)) mock_delete.assert_not_called() - assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected) + assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected) @patch( "exo.download.coordinator.delete_model", @@ -266,12 +264,12 @@ class TestStartDownloadAutoEviction: e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, DownloadEvicted) + and isinstance(e.download_progress, ModelEvicted) 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, DownloadEvicted) + assert isinstance(evicted_dp, ModelEvicted) assert evicted_dp.evicted_for == MODEL_NEW @@ -336,7 +334,7 @@ class TestActiveModelProtection: await _start_download(coordinator, _shard(MODEL_NEW, 5)) mock_delete.assert_not_called() - assert isinstance(coordinator.download_status[MODEL_NEW], DownloadRejected) + assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected) class TestDiskDeleteFailure: @@ -369,7 +367,7 @@ class TestDiskDeleteFailure: # 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) + assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected) # Eviction target should still be in download_status (not removed) assert MODEL_A in coordinator.download_status @@ -449,11 +447,11 @@ class TestEvictionEvents: e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, DownloadEvicted) + and isinstance(e.download_progress, ModelEvicted) ] assert len(evicted_events) == 1 evicted_dp = evicted_events[0].download_progress - assert isinstance(evicted_dp, DownloadEvicted) + assert isinstance(evicted_dp, ModelEvicted) assert evicted_dp.evicted_for == MODEL_NEW assert evicted_dp.shard_metadata.model_card.model_id == MODEL_A @@ -491,7 +489,7 @@ class TestEvictionEvents: e for e in events if isinstance(e, NodeDownloadProgress) - and isinstance(e.download_progress, DownloadEvicted) + and isinstance(e.download_progress, ModelEvicted) ] evicted_model_ids = [ e.download_progress.shard_metadata.model_card.model_id @@ -500,7 +498,7 @@ class TestEvictionEvents: assert evicted_model_ids == [MODEL_A, MODEL_B, MODEL_C] for e in evicted_events: dp = e.download_progress - assert isinstance(dp, DownloadEvicted) + assert isinstance(dp, ModelEvicted) assert dp.evicted_for == MODEL_NEW @@ -512,12 +510,12 @@ class TestClearRejections: config = StorageConfig( max_storage=Memory.from_gb(10), storage_policy="auto-evict" ) - evicted = DownloadEvicted( + evicted = ModelEvicted( node_id=NODE_ID, shard_metadata=_shard(MODEL_A, 4), evicted_for=MODEL_NEW, ) - rejected = DownloadRejected( + rejected = ModelRejected( node_id=NODE_ID, shard_metadata=_shard(MODEL_B, 4), reason="Not enough space", @@ -533,14 +531,14 @@ class TestClearRejections: await coordinator.clear_rejections() # Evicted should remain unchanged - assert isinstance(coordinator.download_status[MODEL_A], DownloadEvicted) + assert isinstance(coordinator.download_status[MODEL_A], ModelEvicted) # Rejected should be cleared to Pending - assert isinstance(coordinator.download_status[MODEL_B], DownloadPending) + assert isinstance(coordinator.download_status[MODEL_B], ModelNotDownloading) async def test_clear_rejections_on_policy_only_change(self) -> None: """clear_rejections fires even when only the policy changes (no limit change).""" config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual") - rejected = DownloadRejected( + rejected = ModelRejected( node_id=NODE_ID, shard_metadata=_shard(MODEL_A, 4), reason="Manual policy", @@ -556,6 +554,6 @@ class TestClearRejections: await coordinator.clear_rejections() # Rejected should be cleared to Pending - assert isinstance(coordinator.download_status[MODEL_A], DownloadPending) + assert isinstance(coordinator.download_status[MODEL_A], ModelNotDownloading) # Completed should be unchanged - assert isinstance(coordinator.download_status[MODEL_B], DownloadCompleted) + assert isinstance(coordinator.download_status[MODEL_B], ModelReady) diff --git a/src/exo/download/tests/test_re_download.py b/src/exo/download/tests/test_re_download.py index 8ecb1187..fda71091 100644 --- a/src/exo/download/tests/test_re_download.py +++ b/src/exo/download/tests/test_re_download.py @@ -21,7 +21,7 @@ from exo.shared.types.commands import ( from exo.shared.types.common import NodeId, SystemId 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.downloads import ModelReady from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata from exo.utils.channels import Receiver, Sender, channel @@ -196,7 +196,7 @@ async def test_re_download_after_delete_completes() -> None: async def _wait_for_download_completed( event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0 -) -> DownloadCompleted | None: +) -> ModelReady | None: """Drain events until we see a DownloadCompleted for the given model, or timeout.""" try: async with asyncio.timeout(timeout): @@ -204,7 +204,7 @@ async def _wait_for_download_completed( event = await event_recv.receive() if ( isinstance(event, NodeDownloadProgress) - and isinstance(event.download_progress, DownloadCompleted) + and isinstance(event.download_progress, ModelReady) and event.download_progress.shard_metadata.model_card.model_id == model_id ): diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 43341650..b07d7e27 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -67,7 +67,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.downloads import ModelRejected from exo.shared.types.worker.instances import InstanceId from exo.utils.channels import Receiver, Sender from exo.utils.disk_event_log import DiskEventLog @@ -430,7 +430,7 @@ class Master: self.state = apply(self.state, indexed) if isinstance(event, NodeDownloadProgress) and isinstance( - event.download_progress, DownloadRejected + event.download_progress, ModelRejected ): dp = event.download_progress rejection_events = get_download_rejected_events( diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index f9380693..c8a7f797 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -32,11 +32,11 @@ from exo.shared.types.memory import Memory from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import ( - DownloadCompleted, - DownloadFailed, - DownloadOngoing, - DownloadPending, - DownloadProgress, + ModelModelDownloadFailed, + ModelDownloading, + ModelNotDownloading, + ModelReady, + ModelStatus, ) from exo.shared.types.worker.instances import ( Instance, @@ -66,26 +66,26 @@ def add_instance_to_placements( def _get_node_download_fraction( node_id: NodeId, model_id: ModelId, - download_status: Mapping[NodeId, Sequence[DownloadProgress]], + download_status: Mapping[NodeId, Sequence[ModelStatus]], ) -> float: """Return the download fraction (0.0–1.0) for a model on a given node.""" for progress in download_status.get(node_id, []): if progress.shard_metadata.model_card.model_id != model_id: continue match progress: - case DownloadCompleted(): + case ModelReady(): return 1.0 - case DownloadOngoing(): + case ModelDownloading(): total = progress.download_progress.total.in_bytes return ( progress.download_progress.downloaded.in_bytes / total if total > 0 else 0.0 ) - case DownloadPending(): + case ModelNotDownloading(): total = progress.total.in_bytes return progress.downloaded.in_bytes / total if total > 0 else 0.0 - case DownloadFailed(): + case ModelDownloadFailed(): return 0.0 return 0.0 @@ -93,7 +93,7 @@ def _get_node_download_fraction( def _cycle_download_score( cycle: Cycle, model_id: ModelId, - download_status: Mapping[NodeId, Sequence[DownloadProgress]], + download_status: Mapping[NodeId, Sequence[ModelStatus]], ) -> float: """Sum of download fractions across all nodes in a cycle.""" return sum( @@ -109,7 +109,7 @@ def place_instance( node_memory: Mapping[NodeId, MemoryUsage], node_network: Mapping[NodeId, NodeNetworkInfo], required_nodes: set[NodeId] | None = None, - download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None, + download_status: Mapping[NodeId, Sequence[ModelStatus]] | None = None, ) -> dict[InstanceId, Instance]: cycles = topology.get_cycles() candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles)) @@ -309,14 +309,14 @@ def get_transition_events( def cancel_unnecessary_downloads( instances: Mapping[InstanceId, Instance], - download_status: Mapping[NodeId, Sequence[DownloadProgress]], + download_status: Mapping[NodeId, Sequence[ModelStatus]], ) -> Sequence[DownloadCommand]: commands: list[DownloadCommand] = [] currently_downloading = [ (k, v.shard_metadata.model_card.model_id) for k, vs in download_status.items() for v in vs - if isinstance(v, (DownloadOngoing)) + if isinstance(v, (ModelDownloading)) ] active_models = set( ( diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index ab6ca4fb..ed83051e 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -40,7 +40,7 @@ from exo.shared.types.profiling import ( from exo.shared.types.state import State 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 +from exo.shared.types.worker.downloads import ModelStatus from exo.shared.types.worker.instances import Instance, InstanceId from exo.shared.types.worker.runners import RunnerId, RunnerShutdown, RunnerStatus from exo.utils.info_gatherer.info_gatherer import ( @@ -131,7 +131,7 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S if not replaced: current.append(dp) - new_downloads: Mapping[NodeId, Sequence[DownloadProgress]] = { + new_downloads: Mapping[NodeId, Sequence[ModelStatus]] = { **state.downloads, node_id: current, } diff --git a/src/exo/shared/storage.py b/src/exo/shared/storage.py index a04bc43e..28214bff 100644 --- a/src/exo/shared/storage.py +++ b/src/exo/shared/storage.py @@ -22,19 +22,19 @@ from exo.shared.types.storage import ( ) from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import ( - DownloadCompleted, - DownloadOngoing, - DownloadProgress, + ModelReady, + ModelDownloading, + ModelStatus, ) from exo.shared.types.worker.instances import Instance, InstanceId -def calculate_used_storage(downloads: Sequence[DownloadProgress]) -> Memory: +def calculate_used_storage(downloads: Sequence[ModelStatus]) -> Memory: total = Memory() for dp in downloads: - if isinstance(dp, DownloadCompleted): + if isinstance(dp, ModelReady): total = total + dp.total - elif isinstance(dp, DownloadOngoing): + elif isinstance(dp, ModelDownloading): total = total + dp.download_progress.total return total @@ -42,7 +42,7 @@ def calculate_used_storage(downloads: Sequence[DownloadProgress]) -> Memory: def check_storage_quota( model_size: Memory, config: StorageConfig, - downloads: Sequence[DownloadProgress], + downloads: Sequence[ModelStatus], ) -> tuple[bool, str]: if config.max_storage is None: return True, "" @@ -60,13 +60,13 @@ def check_storage_quota( def get_lru_eviction_candidates( - downloads: Sequence[DownloadProgress], + downloads: Sequence[ModelStatus], model_last_used: Mapping[ModelId, datetime], active_model_ids: frozenset[ModelId], -) -> list[tuple[ModelId, DownloadCompleted]]: - candidates: list[tuple[ModelId, DownloadCompleted]] = [] +) -> list[tuple[ModelId, ModelReady]]: + candidates: list[tuple[ModelId, ModelReady]] = [] for dp in downloads: - if not isinstance(dp, DownloadCompleted): + if not isinstance(dp, ModelReady): continue if dp.read_only: continue @@ -84,7 +84,7 @@ def get_lru_eviction_candidates( def compute_evictions_needed( model_size: Memory, available: Memory, - candidates: list[tuple[ModelId, DownloadCompleted]], + candidates: list[tuple[ModelId, ModelReady]], ) -> list[ModelId] | None: if model_size <= available: return [] @@ -105,7 +105,7 @@ def compute_evictions_needed( def decide_storage_action( model_size: Memory, config: StorageConfig, - downloads: Sequence[DownloadProgress], + downloads: Sequence[ModelStatus], model_last_used: Mapping[ModelId, datetime], active_model_ids: frozenset[ModelId], ) -> StorageDecision: diff --git a/src/exo/shared/tests/test_apply/test_apply_node_download.py b/src/exo/shared/tests/test_apply/test_apply_node_download.py index f9df6e07..7dc23f72 100644 --- a/src/exo/shared/tests/test_apply/test_apply_node_download.py +++ b/src/exo/shared/tests/test_apply/test_apply_node_download.py @@ -4,14 +4,14 @@ from exo.shared.types.common import NodeId from exo.shared.types.events import NodeDownloadProgress from exo.shared.types.memory import Memory from exo.shared.types.state import State -from exo.shared.types.worker.downloads import DownloadCompleted +from exo.shared.types.worker.downloads import ModelReady from exo.worker.tests.constants import MODEL_A_ID, MODEL_B_ID def test_apply_node_download_progress(): state = State() shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) - event = DownloadCompleted( + event = ModelReady( node_id=NodeId("node-1"), shard_metadata=shard1, total=Memory(), @@ -27,12 +27,12 @@ def test_apply_node_download_progress(): def test_apply_two_node_download_progress(): shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard2 = get_pipeline_shard_metadata(MODEL_B_ID, device_rank=0, world_size=2) - event1 = DownloadCompleted( + event1 = ModelReady( node_id=NodeId("node-1"), shard_metadata=shard1, total=Memory(), ) - event2 = DownloadCompleted( + event2 = ModelReady( node_id=NodeId("node-1"), shard_metadata=shard2, total=Memory(), diff --git a/src/exo/shared/tests/test_storage.py b/src/exo/shared/tests/test_storage.py index 0ae21d50..36409bb4 100644 --- a/src/exo/shared/tests/test_storage.py +++ b/src/exo/shared/tests/test_storage.py @@ -25,9 +25,9 @@ from exo.shared.types.storage import ( ) from exo.shared.types.tasks import LoadModel, TaskId, TaskStatus from exo.shared.types.worker.downloads import ( - DownloadCompleted, - DownloadOngoing, - DownloadPending, + ModelReady, + ModelDownloading, + ModelNotDownloading, DownloadProgressData, ) from exo.shared.types.worker.instances import InstanceId, MlxRingInstance @@ -42,9 +42,9 @@ NODE_ID = "node-1" def _completed( model_id: ModelId, size_gb: float, read_only: bool = False -) -> DownloadCompleted: +) -> ModelReady: shard = get_pipeline_shard_metadata(model_id, device_rank=0) - return DownloadCompleted( + return ModelReady( node_id=NODE_ID, # type: ignore[arg-type] shard_metadata=shard, total=Memory.from_gb(size_gb), @@ -189,11 +189,11 @@ class TestCalculateUsedStorage: shard_b = get_pipeline_shard_metadata(MODEL_B, device_rank=0) downloads = [ _completed(MODEL_A, 5), - DownloadPending( + ModelNotDownloading( node_id=NODE_ID, # type: ignore[arg-type] shard_metadata=shard_a, ), - DownloadOngoing( + ModelDownloading( node_id=NODE_ID, # type: ignore[arg-type] shard_metadata=shard_b, download_progress=DownloadProgressData( @@ -222,7 +222,7 @@ class TestGetLruEvictionCandidatesExtended: shard_a = get_pipeline_shard_metadata(MODEL_A, device_rank=0) downloads = [ _completed(MODEL_B, 3), - DownloadPending( + ModelNotDownloading( node_id=NODE_ID, # type: ignore[arg-type] shard_metadata=shard_a, ), diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index a7aa0ba3..0322db93 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -9,7 +9,7 @@ from exo.shared.types.chunks import GenerationChunk, InputImageChunk from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId from exo.shared.types.storage import StorageConfig from exo.shared.types.tasks import Task, TaskId, TaskStatus -from exo.shared.types.worker.downloads import DownloadProgress +from exo.shared.types.worker.downloads import ModelStatus from exo.shared.types.worker.instances import Instance, InstanceId from exo.shared.types.worker.runners import RunnerId, RunnerStatus from exo.utils.info_gatherer.info_gatherer import GatheredInfo @@ -87,7 +87,7 @@ class NodeGatheredInfo(BaseEvent): class NodeDownloadProgress(BaseEvent): - download_progress: DownloadProgress + download_progress: ModelStatus class ChunkGenerated(BaseEvent): diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py index a9b0ea8b..d0b96d55 100644 --- a/src/exo/shared/types/state.py +++ b/src/exo/shared/types/state.py @@ -19,7 +19,7 @@ from exo.shared.types.profiling import ( ) from exo.shared.types.storage import StorageConfig from exo.shared.types.tasks import Task, TaskId -from exo.shared.types.worker.downloads import DownloadProgress +from exo.shared.types.worker.downloads import ModelStatus from exo.shared.types.worker.instances import Instance, InstanceId from exo.shared.types.worker.runners import RunnerId, RunnerStatus from exo.utils.pydantic_ext import CamelCaseModel @@ -43,7 +43,7 @@ class State(CamelCaseModel): ) instances: Mapping[InstanceId, Instance] = {} runners: Mapping[RunnerId, RunnerStatus] = {} - downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {} + downloads: Mapping[NodeId, Sequence[ModelStatus]] = {} tasks: Mapping[TaskId, Task] = {} last_seen: Mapping[NodeId, datetime] = {} topology: Topology = Field(default_factory=Topology) diff --git a/src/exo/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py index ba1228b8..f08f9b7b 100644 --- a/src/exo/shared/types/worker/downloads.py +++ b/src/exo/shared/types/worker/downloads.py @@ -23,48 +23,48 @@ class DownloadProgressData(CamelCaseModel): files: dict[str, "DownloadProgressData"] -class BaseDownloadProgress(TaggedModel): +class BaseModelStatus(TaggedModel): node_id: NodeId shard_metadata: ShardMetadata model_directory: str = "" -class DownloadPending(BaseDownloadProgress): +class ModelNotDownloading(BaseModelStatus): downloaded: Memory = Memory() total: Memory = Memory() -class DownloadCompleted(BaseDownloadProgress): +class ModelReady(BaseModelStatus): total: Memory read_only: bool = False -class DownloadFailed(BaseDownloadProgress): +class ModelDownloadFailed(BaseModelStatus): error_message: str -class DownloadOngoing(BaseDownloadProgress): +class ModelDownloading(BaseModelStatus): download_progress: DownloadProgressData -class DownloadRejected(BaseDownloadProgress): +class ModelRejected(BaseModelStatus): reason: str required: Memory available: Memory limit: Memory -class DownloadEvicted(BaseDownloadProgress): +class ModelEvicted(BaseModelStatus): evicted_for: ModelId -DownloadProgress = ( - DownloadPending - | DownloadCompleted - | DownloadFailed - | DownloadOngoing - | DownloadRejected - | DownloadEvicted +ModelStatus = ( + ModelNotDownloading + | ModelReady + | ModelDownloadFailed + | ModelDownloading + | ModelRejected + | ModelEvicted ) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index ded0481b..dffad687 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -43,7 +43,7 @@ from exo.shared.types.tasks import ( TextGeneration, ) from exo.shared.types.topology import Connection, SocketConnection -from exo.shared.types.worker.downloads import DownloadCompleted +from exo.shared.types.worker.downloads import ModelReady from exo.shared.types.worker.runners import RunnerId from exo.utils.channels import Receiver, Sender, channel from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer @@ -191,7 +191,7 @@ class Worker: logger.info(f"Model {model_id} found at {found_path}") await self.event_sender.send( NodeDownloadProgress( - download_progress=DownloadCompleted( + download_progress=ModelReady( node_id=self.node_id, shard_metadata=shard, model_directory=str(found_path), diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 0df10ca4..77675bc6 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -20,10 +20,10 @@ from exo.shared.types.tasks import ( TextGeneration, ) from exo.shared.types.worker.downloads import ( - DownloadCompleted, - DownloadFailed, - DownloadOngoing, - DownloadProgress, + ModelReady, + ModelDownloadFailed, + ModelDownloading, + ModelStatus, ) from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId from exo.shared.types.worker.runners import ( @@ -46,7 +46,7 @@ def plan( node_id: NodeId, # Runners is expected to be FRESH and so should not come from state runners: Mapping[RunnerId, RunnerSupervisor], - global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], + global_download_status: Mapping[NodeId, Sequence[ModelStatus]], instances: Mapping[InstanceId, Instance], all_runners: Mapping[RunnerId, RunnerStatus], # all global tasks: Mapping[TaskId, Task], @@ -115,7 +115,7 @@ def _create_runner( def _model_needs_download( node_id: NodeId, runners: Mapping[RunnerId, RunnerSupervisor], - global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], + global_download_status: Mapping[NodeId, Sequence[ModelStatus]], ) -> DownloadModel | None: local_downloads = global_download_status.get(node_id, []) download_status = { @@ -128,7 +128,7 @@ def _model_needs_download( model_id not in download_status or not isinstance( download_status[model_id], - (DownloadOngoing, DownloadCompleted, DownloadFailed), + (ModelDownloading, ModelReady, ModelDownloadFailed), ) ): # We don't invalidate download_status randomly in case a file gets deleted on disk @@ -191,7 +191,7 @@ def _init_distributed_backend( def _load_model( runners: Mapping[RunnerId, RunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], - global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], + global_download_status: Mapping[NodeId, Sequence[ModelStatus]], ) -> LoadModel | None: for runner in runners.values(): instance = runner.bound_instance.instance @@ -200,7 +200,7 @@ def _load_model( all_local_downloads_complete = all( nid in global_download_status and any( - isinstance(dp, DownloadCompleted) + isinstance(dp, ModelReady) and dp.shard_metadata.model_card.model_id == shard_assignments.model_id for dp in global_download_status[nid] ) diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py index abcb4939..d5b018c3 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py +++ b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py @@ -2,7 +2,7 @@ import exo.worker.plan as plan_mod from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.tasks import LoadModel -from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress +from exo.shared.types.worker.downloads import ModelReady, ModelStatus from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import ( RunnerConnected, @@ -89,12 +89,8 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting(): } global_download_status = { - NODE_A: [ - DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) - ], - NODE_B: [ - DownloadCompleted(shard_metadata=shard2, node_id=NODE_B, total=Memory()) - ], + NODE_A: [ModelReady(shard_metadata=shard1, node_id=NODE_A, total=Memory())], + NODE_B: [ModelReady(shard_metadata=shard2, node_id=NODE_B, total=Memory())], } result = plan_mod.plan( @@ -132,10 +128,8 @@ def test_plan_does_not_request_download_when_shard_already_downloaded(): all_runners = {RUNNER_1_ID: RunnerIdle()} # Global state shows shard is downloaded for NODE_A - global_download_status: dict[NodeId, list[DownloadProgress]] = { - NODE_A: [ - DownloadCompleted(shard_metadata=shard, node_id=NODE_A, total=Memory()) - ], + global_download_status: dict[NodeId, list[ModelStatus]] = { + NODE_A: [ModelReady(shard_metadata=shard, node_id=NODE_A, total=Memory())], NODE_B: [], } @@ -180,9 +174,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally(): } global_download_status = { - NODE_A: [ - DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) - ], + NODE_A: [ModelReady(shard_metadata=shard1, node_id=NODE_A, total=Memory())], NODE_B: [], # NODE_B has no downloads completed yet } @@ -198,11 +190,9 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally(): assert result is None global_download_status = { - NODE_A: [ - DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) - ], + NODE_A: [ModelReady(shard_metadata=shard1, node_id=NODE_A, total=Memory())], NODE_B: [ - DownloadCompleted(shard_metadata=shard2, node_id=NODE_B, total=Memory()) + ModelReady(shard_metadata=shard2, node_id=NODE_B, total=Memory()) ], # NODE_B has no downloads completed yet }