diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py index c9336215..6ce4849d 100644 --- a/src/exo/utils/channels.py +++ b/src/exo/utils/channels.py @@ -192,7 +192,13 @@ class MpReceiver[T]: try: return self.receive_nowait() except WouldBlock: - item = self._state.buffer.get() + try: + item = self._state.buffer.get() + except (TypeError, OSError): + # Queue pipe can get closed while we are blocked on get(). + # The underlying connection._handle becomes None, causing + # TypeError in read(handle, remaining). + raise ClosedResourceError from None if isinstance(item, _MpEndOfStream): self.close() raise EndOfStream from None diff --git a/src/exo/utils/info_gatherer/net_profile.py b/src/exo/utils/info_gatherer/net_profile.py index 85986ead..3255b207 100644 --- a/src/exo/utils/info_gatherer/net_profile.py +++ b/src/exo/utils/info_gatherer/net_profile.py @@ -108,7 +108,7 @@ async def check_reachable( await send.send((target_ip, expected_node_id)) async with ( - httpx.AsyncClient(timeout=timeout, limits=limits) as client, + httpx.AsyncClient(timeout=timeout, limits=limits, verify=False) as client, create_task_group() as tg, ): for node_id in topology.list_nodes(): diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index f22ae5c8..47293d8c 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -4,7 +4,7 @@ import resource import time from collections.abc import Generator from functools import cache -from typing import Literal +from typing import TYPE_CHECKING, Literal import mlx.core as mx from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model @@ -588,6 +588,13 @@ def main( case Shutdown(): current_status = RunnerShuttingDown() logger.info("runner shutting down") + if not TYPE_CHECKING: + del inference_model, image_model, tokenizer, group + mx.clear_cache() + import gc + + gc.collect() + event_sender.send( RunnerStatusUpdated( runner_id=runner_id, runner_status=current_status @@ -612,12 +619,8 @@ def main( event_sender.send( RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status) ) - if isinstance(current_status, RunnerShutdown): - del inference_model, image_model, tokenizer, group - mx.clear_cache() - import gc - gc.collect() + if isinstance(current_status, RunnerShutdown): break diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 58ac778e..2cee74fd 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -100,7 +100,6 @@ class RunnerSupervisor: logger.info("Runner supervisor shutting down") self._ev_recv.close() self._task_sender.close() - self._event_sender.close() self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK")) self._cancel_sender.close() self.runner_process.join(5) @@ -180,6 +179,7 @@ class RunnerSupervisor: await self._check_runner(e) for tid in self.pending: self.pending[tid].set() + self._event_sender.close() def __del__(self) -> None: if self.runner_process.is_alive(): @@ -208,10 +208,15 @@ class RunnerSupervisor: logger.opt(exception=e).error(f"Runner terminated ({cause})") - await self._event_sender.send( - RunnerStatusUpdated( - runner_id=self.bound_instance.bound_runner_id, - runner_status=RunnerFailed(error_message=f"Terminated ({cause})"), + try: + await self._event_sender.send( + RunnerStatusUpdated( + runner_id=self.bound_instance.bound_runner_id, + runner_status=RunnerFailed(error_message=f"Terminated ({cause})"), + ) + ) + except (ClosedResourceError, BrokenResourceError): + logger.warning( + "Event sender already closed, unable to report runner failure" ) - ) self.shutdown()