Merge branch 'main' into andrei/force_oom

# Conflicts:
#	uv.lock

.envrc

@@ -1,3 +1,8 @@
+# installs nix-direnv if doesn't exist (speeds up evaluation)
+if ! has nix_direnv_version || ! nix_direnv_version 3.1.1; then
+  source_url "https://raw.githubusercontent.com/nix-community/nix-direnv/3.1.1/direnvrc" "sha256-p+fzQdrms/hDa7g+soShAybJNo4bN4SIAeSfqNKgD5I="
+fi
+
 use flake
 
 # creates .venv if doesn't exist and loads its environment

.idea/misc.xml

@@ -4,4 +4,7 @@
     <option name="sdkName" value="Python 3.13 (exo)" />
   </component>
   <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (exo)" project-jdk-type="Python SDK" />
+  <component name="RuffConfiguration">
+    <option name="enabled" value="true" />
+  </component>
 </project>

src/exo/shared/types/chunks.py

@@ -13,6 +13,7 @@ from exo.shared.models.model_cards import ModelId
 from exo.utils.pydantic_ext import TaggedModel
 from exo.worker.runner.diagnostics import KnownRunnerDiagnostic
 
+from ...worker.runner.diagnostics import KnownRunnerDiagnostic
 from .common import CommandId
 
 

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/worker/engines/mlx/utils_mlx.py

@@ -1,3 +1,4 @@
+from itertools import pairwise
 import json
 import os
 import re
@@ -5,6 +6,7 @@ import sys
 import tempfile
 import time
 from collections.abc import Generator
+from math import isqrt
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
@@ -786,13 +788,75 @@ def mlx_force_oom(size: int = 200000) -> None:
     a = mx.random.uniform(shape=(size, size), dtype=mx.float32)
     b = mx.random.uniform(shape=(size, size), dtype=mx.float32)
     mx.eval(a, b)
-    c = mx.matmul(a, b)
-    d = mx.matmul(a, c)
-    e = mx.matmul(b, c)
-    f = mx.sigmoid(d + e)
+    c = mx.matmul(a, b)  # (size,size)
+    d = mx.matmul(a, c)  # (size,size)
+    e = mx.matmul(b, c)  # (size,size)
+    f = mx.sigmoid(d + e)  # (size,size)
     mx.eval(f)
 
 
+def mlx_force_oom2(bytes_alloc: int = 1024**5):  # the default is 1 petabyte lol
+    """
+    Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
+
+    NOTE: probably only works correctly on Apple unified memory
+    """
+
+    # TODO: I give up. this either uses swap (inconsistently) and only sometimes OOMs
+    #       or if I tune it to be more aggressive then it kerenel panics entirely;
+    #       there MIGHT be a way to make it not use swap memory but I'm not determined enough
+    #       to figure out how :)
+
+    def get_size(memory: int):
+        mat_elem = -(-memory // 4)  # per-matrix elements (4 bytes per elem)
+
+        # square root to get size (round up if not integer)
+        root = isqrt(mat_elem)
+        return root if root**2 == mat_elem else root + 1
+
+    def oom(size: int):
+        mx.set_default_device(mx.gpu)
+        mx.clear_cache()
+
+        # allocate a lot
+        z = mx.zeros(shape=(size, size), dtype=mx.float32)
+        t1 = [mx.random.uniform(shape=(size, size), dtype=mx.float32) for _ in range(2)]
+
+        # mat mul cycle
+        t2: list[mx.array] = []
+        for m1, m2 in pairwise(t1):
+            t2.append(mx.matmul(m1, m2))
+        print("t2-run")
+        mx.eval(*t2)
+        print("t2-eval")
+
+        # sigmoid sum
+        f = mx.sigmoid(sum(t1, start=z) + sum(t2, start=z))
+        print("f-run")
+        mx.eval(f)
+        print("f-eval")
+
+    # use supplied size, or computer appropriate size otherwise
+    fail_num = 0
+    while True:
+        try:
+            print(f"size {bytes_alloc / 1024**3} GB")
+            oom(get_size(bytes_alloc))
+            break
+        except RuntimeError as e:
+            max_bytes = re.compile(
+                r"\[metal::malloc\] Attempting to allocate (?:\d+) bytes which is greater than the maximum allowed buffer size of (?P<max_bytes>\d+) bytes."
+            ).match(str(e))
+            if max_bytes is None:
+                raise RuntimeError(
+                    "Tried to get max buffer, but wrong error format"
+                ) from e
+            bytes_alloc = round(int(max_bytes.group("max_bytes")) * 0.95**fail_num)
+            fail_num += 1
+
+    mlx_force_oom2()
+
+
 def set_wired_limit_for_model(model_size: Memory):
     """
     A context manager to temporarily change the wired limit.