diff --git a/src/exo/utils/info_gatherer/macmon.py b/src/exo/utils/info_gatherer/macmon.py index 4fe714e27..3dc25bb60 100644 --- a/src/exo/utils/info_gatherer/macmon.py +++ b/src/exo/utils/info_gatherer/macmon.py @@ -1,3 +1,6 @@ +import os +import shutil +import subprocess from typing import Self from pydantic import BaseModel @@ -68,3 +71,41 @@ class MacmonMetrics(TaggedModel): @classmethod def from_raw_json(cls, json: str) -> Self: return cls.from_raw(RawMacmonMetrics.model_validate_json(json)) + + +def read_macmon_metrics_once( + macmon_path: str | None = None, + *, + timeout: float = 5, +) -> MacmonMetrics | None: + """ + Read a single macmon sample, returning None when macmon is unavailable. + """ + resolved_macmon_path = ( + macmon_path or os.getenv("EXO_MACMON_PATH") or shutil.which("macmon") + ) + if resolved_macmon_path is None: + return None + + try: + result = subprocess.run( + [resolved_macmon_path, "pipe", "--samples", "1", "--interval", "100"], + capture_output=True, + check=False, + text=True, + timeout=timeout, + ) + except (OSError, subprocess.SubprocessError): + return None + + if result.returncode != 0: + return None + + lines = result.stdout.strip().splitlines() + if not lines: + return None + + try: + return MacmonMetrics.from_raw_json(lines[0]) + except ValueError: + return None diff --git a/src/exo/utils/info_gatherer/tests/test_macmon.py b/src/exo/utils/info_gatherer/tests/test_macmon.py new file mode 100644 index 000000000..6ce3c485d --- /dev/null +++ b/src/exo/utils/info_gatherer/tests/test_macmon.py @@ -0,0 +1,71 @@ +import subprocess + +import pytest + +import exo.utils.info_gatherer.macmon as macmon +from exo.utils.info_gatherer.macmon import read_macmon_metrics_once + + +def _macmon_output() -> str: + return """ +{"timestamp":"2026-05-14T12:00:00Z","temp":{"cpu_temp_avg":45.0,"gpu_temp_avg":46.0},"memory":{"ram_total":1000,"ram_usage":275,"swap_total":500,"swap_usage":125},"ecpu_usage":[1000,1.0],"pcpu_usage":[2000,2.0],"gpu_usage":[1200,3.0],"all_power":1.0,"ane_power":2.0,"cpu_power":3.0,"gpu_power":4.0,"gpu_ram_power":5.0,"ram_power":6.0,"sys_power":7.0} +{"timestamp":"2026-05-14T12:00:01Z","temp":{"cpu_temp_avg":45.0,"gpu_temp_avg":46.0},"memory":{"ram_total":1000,"ram_usage":999,"swap_total":500,"swap_usage":125},"ecpu_usage":[1000,1.0],"pcpu_usage":[2000,2.0],"gpu_usage":[1200,3.0],"all_power":1.0,"ane_power":2.0,"cpu_power":3.0,"gpu_power":4.0,"gpu_ram_power":5.0,"ram_power":6.0,"sys_power":7.0} +""" + + +def test_read_macmon_metrics_once_uses_first_sample( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[list[str]] = [] + + def fake_run( + args: list[str], + *, + capture_output: bool, + check: bool, + text: bool, + timeout: float, + ) -> subprocess.CompletedProcess[str]: + calls.append(args) + assert capture_output + assert not check + assert text + assert timeout == 3 + return subprocess.CompletedProcess( + args=args, + returncode=0, + stdout=_macmon_output(), + stderr="", + ) + + monkeypatch.setattr(macmon.subprocess, "run", fake_run) + + metrics = read_macmon_metrics_once("/usr/local/bin/macmon", timeout=3) + + assert metrics is not None + assert metrics.memory.ram_available.in_bytes == 725 + assert calls == [ + ["/usr/local/bin/macmon", "pipe", "--samples", "1", "--interval", "100"] + ] + + +def test_read_macmon_metrics_once_returns_none_for_empty_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fake_run( + args: list[str], + *, + capture_output: bool, + check: bool, + text: bool, + timeout: float, + ) -> subprocess.CompletedProcess[str]: + return subprocess.CompletedProcess(args=args, returncode=0, stdout="", stderr="") + + monkeypatch.setattr(macmon.subprocess, "run", fake_run) + + assert read_macmon_metrics_once("/usr/local/bin/macmon") is None + + +def test_read_macmon_metrics_once_returns_none_for_missing_binary() -> None: + assert read_macmon_metrics_once("/does/not/exist/macmon") is None diff --git a/src/exo/worker/engines/mlx/tests/test_utils_mlx.py b/src/exo/worker/engines/mlx/tests/test_utils_mlx.py index 0dbcf71ec..d7b9f4e75 100644 --- a/src/exo/worker/engines/mlx/tests/test_utils_mlx.py +++ b/src/exo/worker/engines/mlx/tests/test_utils_mlx.py @@ -1,11 +1,6 @@ -# pyright: reportPrivateUsage=false - import pytest -from exo.worker.engines.mlx.utils_mlx import ( - _available_memory_from_macmon_output, - get_mlx_force_oom_size, -) +from exo.worker.engines.mlx.utils_mlx import get_mlx_force_oom_size _MLX_FORCE_OOM_DTYPE_BYTES = 4 _MLX_FORCE_OOM_LIVE_MATRICES = 3 @@ -43,16 +38,3 @@ def test_get_mlx_force_oom_size_requires_positive_available_ram( ) -> None: with pytest.raises(ValueError, match="available_ram must be positive"): get_mlx_force_oom_size(available_ram) - - -def test_available_memory_from_macmon_output_uses_first_sample() -> None: - output = """ -{"timestamp":"2026-05-14T12:00:00Z","temp":{"cpu_temp_avg":45.0,"gpu_temp_avg":46.0},"memory":{"ram_total":1000,"ram_usage":275,"swap_total":500,"swap_usage":125},"ecpu_usage":[1000,1.0],"pcpu_usage":[2000,2.0],"gpu_usage":[1200,3.0],"all_power":1.0,"ane_power":2.0,"cpu_power":3.0,"gpu_power":4.0,"gpu_ram_power":5.0,"ram_power":6.0,"sys_power":7.0} -{"timestamp":"2026-05-14T12:00:01Z","temp":{"cpu_temp_avg":45.0,"gpu_temp_avg":46.0},"memory":{"ram_total":1000,"ram_usage":999,"swap_total":500,"swap_usage":125},"ecpu_usage":[1000,1.0],"pcpu_usage":[2000,2.0],"gpu_usage":[1200,3.0],"all_power":1.0,"ane_power":2.0,"cpu_power":3.0,"gpu_power":4.0,"gpu_ram_power":5.0,"ram_power":6.0,"sys_power":7.0} -""" - - assert _available_memory_from_macmon_output(output) == 725 - - -def test_available_memory_from_macmon_output_rejects_missing_sample() -> None: - assert _available_memory_from_macmon_output("") is None diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 4bc4feead..b4dfb3428 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -1,8 +1,6 @@ import json import os import re -import shutil -import subprocess import sys import tempfile import time @@ -61,7 +59,7 @@ from exo.shared.types.worker.shards import ( ShardMetadata, TensorShardMetadata, ) -from exo.utils.info_gatherer.macmon import MacmonMetrics +from exo.utils.info_gatherer.macmon import read_macmon_metrics_once from exo.worker.engines.mlx.auto_parallel import ( get_inner_model, get_layers, @@ -832,16 +830,6 @@ def get_mlx_force_oom_size(available_ram: int) -> int: return root + 1 -def _available_memory_from_macmon_output(output: str) -> int | None: - lines = output.strip().splitlines() - if not lines: - return None - try: - return MacmonMetrics.from_raw_json(lines[0]).memory.ram_available.in_bytes - except ValueError: - return None - - def mlx_force_oom2() -> None: """ Force an MLX Metal OOM using the current machine's available unified memory. @@ -851,23 +839,9 @@ def mlx_force_oom2() -> None: available_memory: int | None = None if sys.platform == "darwin": - macmon_path = os.getenv("EXO_MACMON_PATH") or shutil.which("macmon") - if macmon_path is not None: - try: - result = subprocess.run( - [macmon_path, "pipe", "--samples", "1", "--interval", "100"], - capture_output=True, - check=False, - text=True, - timeout=5, - ) - except (OSError, subprocess.SubprocessError): - result = None - - if result is not None and result.returncode == 0: - # macmon reports unified RAM in use more accurately on Apple - # Silicon; subtract usage from total via the parser above. - available_memory = _available_memory_from_macmon_output(result.stdout) + macmon_metrics = read_macmon_metrics_once() + if macmon_metrics is not None: + available_memory = macmon_metrics.memory.ram_available.in_bytes if available_memory is None: import psutil