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

Evict directly instead of through DownloadRejected

Browse Source
This commit is contained in:
ciaranbor
2026-03-10 10:40:22 +00:00
parent e1ddaf97bd
commit 9832f161e7
4 changed files with 151 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/download/coordinator.py
+10 -20
View File
@@ -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
src/exo/download/tests/test_auto_eviction.py
+131 -5
View File
@@ -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)
src/exo/main.py
+9 -5
View File
@@ -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()
src/exo/utils/info_gatherer/info_gatherer.py
+1 -1
View File
@@ -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,