diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py index de9f4e97..a79f9f70 100644 --- a/src/exo/utils/info_gatherer/info_gatherer.py +++ b/src/exo/utils/info_gatherer/info_gatherer.py @@ -371,20 +371,8 @@ GatheredInfo = ( @dataclass class InfoGatherer: info_sender: Sender[GatheredInfo] - interface_watcher_interval: float | None = 10 - misc_poll_interval: float | None = 60 - system_profiler_interval: float | None = 5 if IS_DARWIN else None - memory_poll_rate: float | None = None if IS_DARWIN else 1 - macmon_interval: float | None = 1 if IS_DARWIN else None - thunderbolt_bridge_poll_interval: float | None = 10 if IS_DARWIN else None - static_info_poll_interval: float | None = 60 - rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None - disk_poll_interval: float | None = 30 _tg: TaskGroup = field(init=False, default_factory=TaskGroup) - _psutil_memory_fallback_enabled: bool = field(init=False, default=False) - - def _get_macmon_path(self) -> str | None: - return os.getenv("EXO_MACMON_PATH") or shutil.which("macmon") + _psutil_enabled: bool = field(init=False, default=False) async def _can_read_macmon_metrics(self, macmon_path: str) -> bool: try: @@ -425,25 +413,16 @@ class InfoGatherer: async def run(self): async with self._tg as tg: if IS_DARWIN: - if (macmon_path := self._get_macmon_path()) is not None: - if await self._can_read_macmon_metrics(macmon_path): - tg.start_soon(self._monitor_macmon, macmon_path) - else: - logger.warning( - f"macmon at {macmon_path} is unusable, falling back to psutil memory monitoring" - ) - else: - logger.warning( - "macmon not found, falling back to psutil for memory monitoring" - ) - tg.start_soon(self._monitor_system_profiler_thunderbolt_data) - tg.start_soon(self._monitor_thunderbolt_bridge_status) - tg.start_soon(self._monitor_rdma_ctl_status) - tg.start_soon(self._watch_system_info) - tg.start_soon(self._monitor_memory_usage) - tg.start_soon(self._monitor_misc) - tg.start_soon(self._monitor_static_info) - tg.start_soon(self._monitor_disk_usage) + tg.start_soon(self._monitor_macmon, 1) + tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5) + tg.start_soon(self._monitor_thunderbolt_bridge_status, 10) + tg.start_soon(self._monitor_rdma_ctl_status, 10) + if not IS_DARWIN: + tg.start_soon(self._monitor_memory_usage, 1) + tg.start_soon(self._watch_system_info, 10) + tg.start_soon(self._monitor_misc, 60) + tg.start_soon(self._monitor_static_info, 60) + tg.start_soon(self._monitor_disk_usage, 30) nc = await NodeConfig.gather() if nc is not None: @@ -452,32 +431,27 @@ class InfoGatherer: def shutdown(self): self._tg.cancel_tasks() - async def _monitor_static_info(self): - if self.static_info_poll_interval is None: - return + async def _monitor_static_info(self, static_info_poll_interval: float): while True: try: with fail_after(30): await self.info_sender.send(await StaticNodeInformation.gather()) except Exception as e: logger.opt(exception=e).warning("Error gathering static node info") - await anyio.sleep(self.static_info_poll_interval) + await anyio.sleep(static_info_poll_interval) - async def _monitor_misc(self): - if self.misc_poll_interval is None: - return + async def _monitor_misc(self, misc_poll_interval: float): while True: try: with fail_after(10): await self.info_sender.send(await MiscData.gather()) except Exception as e: logger.opt(exception=e).warning("Error gathering misc data") - await anyio.sleep(self.misc_poll_interval) - - async def _monitor_system_profiler_thunderbolt_data(self): - if self.system_profiler_interval is None: - return + await anyio.sleep(misc_poll_interval) + async def _monitor_system_profiler_thunderbolt_data( + self, system_profiler_interval: float + ): while True: try: with fail_after(30): @@ -499,17 +473,18 @@ class InfoGatherer: await self.info_sender.send(MacThunderboltConnections(conns=conns)) except Exception as e: logger.opt(exception=e).warning("Error gathering Thunderbolt data") - await anyio.sleep(self.system_profiler_interval) + await anyio.sleep(system_profiler_interval) - async def _monitor_memory_usage(self): + async def _monitor_memory_usage(self, memory_poll_rate: float): + if self._psutil_enabled: + return + self._psutil_enabled = True override_memory_env = os.getenv("OVERRIDE_MEMORY_MB") override_memory: int | None = ( Memory.from_mb(int(override_memory_env)).in_bytes if override_memory_env else None ) - if self.memory_poll_rate is None: - return while True: try: await self.info_sender.send( @@ -517,11 +492,9 @@ class InfoGatherer: ) except Exception as e: logger.opt(exception=e).warning("Error gathering memory usage") - await anyio.sleep(self.memory_poll_rate) + await anyio.sleep(memory_poll_rate) - async def _watch_system_info(self): - if self.interface_watcher_interval is None: - return + async def _watch_system_info(self, interface_watcher_interval: float): while True: try: with fail_after(10): @@ -529,11 +502,11 @@ class InfoGatherer: await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics)) except Exception as e: logger.opt(exception=e).warning("Error gathering network interfaces") - await anyio.sleep(self.interface_watcher_interval) + await anyio.sleep(interface_watcher_interval) - async def _monitor_thunderbolt_bridge_status(self): - if self.thunderbolt_bridge_poll_interval is None: - return + async def _monitor_thunderbolt_bridge_status( + self, thunderbolt_bridge_poll_interval: float + ): while True: try: with fail_after(30): @@ -544,11 +517,9 @@ class InfoGatherer: logger.opt(exception=e).warning( "Error gathering Thunderbolt Bridge status" ) - await anyio.sleep(self.thunderbolt_bridge_poll_interval) + await anyio.sleep(thunderbolt_bridge_poll_interval) - async def _monitor_rdma_ctl_status(self): - if self.rdma_ctl_poll_interval is None: - return + async def _monitor_rdma_ctl_status(self, rdma_ctl_poll_interval: float): while True: try: curr = await RdmaCtlStatus.gather() @@ -556,26 +527,36 @@ class InfoGatherer: await self.info_sender.send(curr) except Exception as e: logger.opt(exception=e).warning("Error gathering RDMA ctl status") - await anyio.sleep(self.rdma_ctl_poll_interval) + await anyio.sleep(rdma_ctl_poll_interval) - async def _monitor_disk_usage(self): - if self.disk_poll_interval is None: - return + async def _monitor_disk_usage(self, disk_poll_interval: float): while True: try: with fail_after(5): await self.info_sender.send(await NodeDiskUsage.gather()) except Exception as e: logger.opt(exception=e).warning("Error gathering disk usage") - await anyio.sleep(self.disk_poll_interval) + await anyio.sleep(disk_poll_interval) - async def _monitor_macmon(self, macmon_path: str): - if self.macmon_interval is None: + async def _monitor_macmon(self, macmon_interval: float): + if ( + macmon_path := os.getenv("EXO_MACMON_PATH") or shutil.which("macmon") + ) is None: + logger.warning( + "macmon not found, falling back to psutil for memory monitoring" + ) + self._tg.start_soon(self._monitor_memory_usage, 1) + return + if not await self._can_read_macmon_metrics(macmon_path): + logger.warning( + f"macmon at {macmon_path} is unusable, falling back to psutil memory monitoring" + ) + self._tg.start_soon(self._monitor_memory_usage, 1) return # macmon pipe --interval [interval in ms] # Timeout: if macmon produces no output for this many seconds, restart it. # macmon writes every macmon_interval seconds, so 10x that is generous. - read_timeout = max(self.macmon_interval * 10, 30) + read_timeout = max(macmon_interval * 10, 30) while True: try: async with await open_process( @@ -583,7 +564,7 @@ class InfoGatherer: macmon_path, "pipe", "--interval", - str(self.macmon_interval * 1000), + str(macmon_interval * 1000), ] ) as p: if not p.stdout: @@ -602,8 +583,7 @@ class InfoGatherer: logger.warning( f"MacMon produced no output for {read_timeout}s, restarting" ) - self.memory_poll_rate = 1 - self._tg.start_soon(self._monitor_memory_usage) + self._tg.start_soon(self._monitor_memory_usage, 1) except CalledProcessError as e: stderr_msg = "no stderr" stderr_output = cast(bytes | str | None, e.stderr) @@ -616,10 +596,8 @@ class InfoGatherer: logger.warning( f"MacMon failed with return code {e.returncode}: {stderr_msg}" ) - self.memory_poll_rate = 1 - self._tg.start_soon(self._monitor_memory_usage) + self._tg.start_soon(self._monitor_memory_usage, 1) except Exception as e: logger.opt(exception=e).warning("Error in macmon monitor") - self.memory_poll_rate = 1 - self._tg.start_soon(self._monitor_memory_usage) - await anyio.sleep(self.macmon_interval) + self._tg.start_soon(self._monitor_memory_usage, 1) + await anyio.sleep(macmon_interval)