Move one-shot macmon sampling helper

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

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

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

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