From fa29d0fd00e8f57830a0c8b25042887b8ab634db Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Mon, 9 Mar 2026 18:59:47 +0000 Subject: [PATCH] Deduplicate by model ID --- src/exo/download/tests/test_auto_eviction.py | 5 +++-- src/exo/shared/apply.py | 9 ++------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src/exo/download/tests/test_auto_eviction.py b/src/exo/download/tests/test_auto_eviction.py index 2512be4c..82032d86 100644 --- a/src/exo/download/tests/test_auto_eviction.py +++ b/src/exo/download/tests/test_auto_eviction.py @@ -129,8 +129,9 @@ class TestStartDownloadAutoEviction: # MODEL_A (oldest) should have been evicted mock_delete.assert_called_once_with(MODEL_A) - assert isinstance(coordinator.download_status[MODEL_A], DownloadEvicted) - assert coordinator.download_status[MODEL_A].evicted_for == MODEL_NEW + evicted_status = coordinator.download_status[MODEL_A] + assert isinstance(evicted_status, DownloadEvicted) + assert evicted_status.evicted_for == MODEL_NEW @patch( "exo.download.coordinator.delete_model", diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 920e1336..3d059890 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -119,18 +119,13 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S """ dp = event.download_progress node_id = dp.node_id + model_id = dp.shard_metadata.model_card.model_id current = list(state.downloads.get(node_id, ())) replaced = False for i, existing_dp in enumerate(current): - # TODO(ciaran): deduplicate by model_id for now. Will need to use - # shard_metadata again when pipeline and tensor downloads differ. - # For now this is fine - if ( - existing_dp.shard_metadata.model_card.model_id - == dp.shard_metadata.model_card.model_id - ): + if existing_dp.shard_metadata.model_card.model_id == model_id: current[i] = dp replaced = True break