fix: garbage-collect orphaned download state

src/exo/master/main.py

@@ -79,6 +79,12 @@
 from exo.utils.task_group import TaskGroup
 
 
+def orphaned_download_node_ids(state: State) -> set[NodeId]:
+    """Return download owners that are no longer members of the topology."""
+    connected_node_ids = set(state.topology.list_nodes())
+    return set(state.downloads).difference(connected_node_ids)
+
+
 def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
     decode = state.instances.get(decode_instance_id)
     if decode is None:
@@ -470,6 +476,12 @@ async def _command_processor(self) -> None:
     # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
     async def _plan(self) -> None:
         while True:
+            # Garbage-collect downloads left behind by node IDs that disappeared
+            # before the master observed a last-seen timeout for them.
+            for node_id in orphaned_download_node_ids(self.state):
+                logger.info(f"Removing downloads belonging to retired node {node_id}")
+                await self.event_sender.send(NodeTimedOut(node_id=node_id))
+
             # kill broken instances
             connected_node_ids = set(self.state.topology.list_nodes())
             for instance_id, instance in self.state.instances.items():

src/exo/master/tests/test_orphaned_downloads.py

@@ -0,0 +1,27 @@
+from exo.master.main import orphaned_download_node_ids
+from exo.shared.tests.conftest import get_pipeline_shard_metadata
+from exo.shared.topology import Topology
+from exo.shared.types.common import NodeId
+from exo.shared.types.state import State
+from exo.shared.types.worker.downloads import DownloadPending
+from exo.worker.tests.constants import MODEL_A_ID
+
+
+def test_orphaned_download_node_ids_excludes_connected_nodes() -> None:
+    retired_node = NodeId("retired-node")
+    connected_node = NodeId("connected-node")
+    topology = Topology()
+    topology.add_node(connected_node)
+    shard = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=1)
+
+    state = State(
+        topology=topology,
+        downloads={
+            retired_node: [DownloadPending(node_id=retired_node, shard_metadata=shard)],
+            connected_node: [
+                DownloadPending(node_id=connected_node, shard_metadata=shard)
+            ],
+        },
+    )
+
+    assert orphaned_download_node_ids(state) == {retired_node}

src/exo/shared/tests/test_apply/test_apply_node_download.py

@@ -1,10 +1,10 @@
-from exo.shared.apply import apply_node_download_progress
+from exo.shared.apply import apply_node_download_progress, apply_node_timed_out
 from exo.shared.tests.conftest import get_pipeline_shard_metadata
 from exo.shared.types.common import NodeId
-from exo.shared.types.events import NodeDownloadProgress
+from exo.shared.types.events import NodeDownloadProgress, NodeTimedOut
 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 DownloadCompleted, DownloadPending
 from exo.worker.tests.constants import MODEL_A_ID, MODEL_B_ID
 
 
@@ -44,3 +44,27 @@ def test_apply_two_node_download_progress():
     )
 
     assert new_state.downloads == {NodeId("node-1"): [event1, event2]}
+
+
+def test_apply_node_timed_out_removes_pending_downloads() -> None:
+    retired_node = NodeId("retired-node")
+    connected_node = NodeId("connected-node")
+    shard = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=1)
+    retired_download = DownloadPending(
+        node_id=retired_node,
+        shard_metadata=shard,
+    )
+    connected_download = DownloadPending(
+        node_id=connected_node,
+        shard_metadata=shard,
+    )
+    state = State(
+        downloads={
+            retired_node: [retired_download],
+            connected_node: [connected_download],
+        }
+    )
+
+    new_state = apply_node_timed_out(NodeTimedOut(node_id=retired_node), state)
+
+    assert new_state.downloads == {connected_node: [connected_download]}