From 4ccf13f9304d0d7ab4b8c2e5f58f0193bed69721 Mon Sep 17 00:00:00 2001 From: Jake Abendroth Date: Thu, 9 Jul 2026 23:13:46 -0700 Subject: [PATCH 1/2] fix: garbage-collect orphaned download state --- src/exo/master/main.py | 12 ++++++++ .../master/tests/test_orphaned_downloads.py | 29 ++++++++++++++++++ .../test_apply/test_apply_node_download.py | 30 +++++++++++++++++-- 3 files changed, 68 insertions(+), 3 deletions(-) create mode 100644 src/exo/master/tests/test_orphaned_downloads.py diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 485ede30f7..9a4630e1c4 100644 --- a/src/exo/master/main.py +++ b/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(): diff --git a/src/exo/master/tests/test_orphaned_downloads.py b/src/exo/master/tests/test_orphaned_downloads.py new file mode 100644 index 0000000000..1b26ad8bac --- /dev/null +++ b/src/exo/master/tests/test_orphaned_downloads.py @@ -0,0 +1,29 @@ +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} 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 f9df6e07a0..ffead5def6 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 @@ -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]} From f98faee85b29d9b917e457f20b98d315a7a6bbb5 Mon Sep 17 00:00:00 2001 From: Jake Abendroth Date: Fri, 10 Jul 2026 01:01:52 -0700 Subject: [PATCH 2/2] style: apply ruff formatting --- src/exo/master/tests/test_orphaned_downloads.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/exo/master/tests/test_orphaned_downloads.py b/src/exo/master/tests/test_orphaned_downloads.py index 1b26ad8bac..91936d1e22 100644 --- a/src/exo/master/tests/test_orphaned_downloads.py +++ b/src/exo/master/tests/test_orphaned_downloads.py @@ -17,9 +17,7 @@ def test_orphaned_download_node_ids_excludes_connected_nodes() -> None: state = State( topology=topology, downloads={ - retired_node: [ - DownloadPending(node_id=retired_node, shard_metadata=shard) - ], + retired_node: [DownloadPending(node_id=retired_node, shard_metadata=shard)], connected_node: [ DownloadPending(node_id=connected_node, shard_metadata=shard) ],