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..91936d1e22 --- /dev/null +++ b/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} 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]}