diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 5c55970f..06ff54ee 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -49,6 +49,7 @@ class DownloadCoordinator: active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict) _tg: TaskGroup = field(init=False, default_factory=TaskGroup) + _stopped: anyio.Event = field(init=False, default_factory=anyio.Event) # Per-model throttle for download progress events _last_progress_time: dict[ModelId, float] = field(default_factory=dict) @@ -100,12 +101,16 @@ class DownloadCoordinator: logger.info( f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}" ) - async with self._tg as tg: - tg.start_soon(self._command_processor) - tg.start_soon(self._emit_existing_download_progress) + try: + async with self._tg as tg: + tg.start_soon(self._command_processor) + tg.start_soon(self._emit_existing_download_progress) + finally: + self._stopped.set() - def shutdown(self) -> None: + async def shutdown(self) -> None: self._tg.cancel_tasks() + await self._stopped.wait() async def _command_processor(self) -> None: with self.download_command_receiver as commands: diff --git a/src/exo/download/tests/test_re_download.py b/src/exo/download/tests/test_re_download.py index 3f159814..75617997 100644 --- a/src/exo/download/tests/test_re_download.py +++ b/src/exo/download/tests/test_re_download.py @@ -186,7 +186,7 @@ async def test_re_download_after_delete_completes() -> None: "Re-download after deletion should complete" ) finally: - coordinator.shutdown() + await coordinator.shutdown() coordinator_task.cancel() with contextlib.suppress(asyncio.CancelledError): await coordinator_task diff --git a/src/exo/main.py b/src/exo/main.py index 93922cc2..68cd70d4 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -228,7 +228,7 @@ class Node: ) if result.is_new_master: if self.download_coordinator: - self.download_coordinator.shutdown() + await self.download_coordinator.shutdown() self.download_coordinator = DownloadCoordinator( self.node_id, exo_shard_downloader(offline=self.offline), @@ -240,7 +240,7 @@ class Node: ) self._tg.start_soon(self.download_coordinator.run) if self.worker: - self.worker.shutdown() + await self.worker.shutdown() # TODO: add profiling etc to resource monitor self.worker = Worker( self.node_id, diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 4bfed825..833f9134 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -80,6 +80,7 @@ class Worker: self.input_chunk_counts: dict[CommandId, int] = {} self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0) + self._stopped: anyio.Event = anyio.Event() async def run(self): logger.info("Starting Worker") @@ -102,6 +103,7 @@ class Worker: self.download_command_sender.close() for runner in self.runners.values(): runner.shutdown() + self._stopped.set() async def _forward_info(self, recv: Receiver[GatheredInfo]): with recv as info_stream: @@ -280,8 +282,9 @@ class Worker: case task: await self._start_runner_task(task) - def shutdown(self): + async def shutdown(self): self._tg.cancel_tasks() + await self._stopped.wait() async def _start_runner_task(self, task: Task): if (instance := self.state.instances.get(task.instance_id)) is not None: