feat: configure GPU wired memory before starting oMLX

README.md

@@ -157,6 +157,14 @@ When `services.olla.version` is omitted inside an explicit `services.olla` block
 
 After changing model placement or node topology in `tfconfig.yaml`, run `make restart gateway-cache-01` (or full `make restart`) before `make smoke <node>` so Olla and TF edge reload the generated router config.
 
+### GPU memory ceiling on inference nodes
+
+Set `services.omlx.gpu_memory_gb` in `tfconfig.yaml` for a shared ceiling, or override it with `nodes.<node>.runtime.gpu_memory_gb`. Use a positive integer in **GiB**. Omitted node values inherit the service setting; explicit `null` selects actual RAM minus 10 GiB. A 128 GiB node therefore defaults to 118 GiB (`iogpu.wired_limit_mb=120832`). RAM comes from `hw.memsize` on the inference node, not the controller or the inventory's `ram_gb`.
+
+`make bootstrap <node>`, `make restart <node>`, and sync with runtime restart install the setting in the oMLX system LaunchDaemon. At every start, including after reboot, a short root-owned plist command validates the limit, applies it only when changed, verifies it by reading sysctl back, then uses `sudo -n -H -u <node.user>` to start oMLX as the non-root runtime user. No extra sysctl sudoers permission is granted. An unsupported host, invalid limit, or failed write/readback prevents oMLX startup and reports the error in its stderr log. Dry-run only renders the plist. The legacy user LaunchAgent and process managers do not apply GPU memory settings; use the system `daemon` manager.
+
+This is a **host-wide ceiling**, shared by all GPU processes, not a reserved allocation or an Olla router option. It must be strictly below physical RAM. It does not override oMLX's memory guard (`--memory-guard` or `--memory-guard-gb`); review that separately if its ceiling remains lower. Raising the ceiling can increase host memory pressure. To roll back, set the previous explicit GiB limit and restart; for the macOS automatic policy, remove the startup policy before resetting `sudo sysctl -w iogpu.wired_limit_mb=0`, otherwise the next daemon start reapplies it. [MLX wired-limit documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.set_wired_limit.html).
+
 Bootstrap verifies minimal service readiness only: Olla `/internal/health`, TF edge auth boundary health, and direct oMLX `/health`. It does not require inference models or chat to be ready; use `make smoke` for model visibility, routing, and chat checks after the cluster has warmed up.
 
 Run system-daemon install targets from a real terminal, not the VS Code guarded terminal, because macOS sudo/su password prompts can be blocked by the editor guard:
@@ -210,6 +218,7 @@ Thunder Forge service ports live in `tfconfig.yaml` under `services:`:
 | `services.edge.host` | `0.0.0.0` | TF edge bind address; use `0.0.0.0` for LAN clients, keep raw Olla private |
 | `services.edge.port` | `40116` | Local TF edge OpenAI-compatible proxy |
 | `services.omlx.port` | `8018` | Default oMLX node runtime port when a node runtime omits `port` |
+| `services.omlx.gpu_memory_gb` | `null` | Host-wide GPU wired-memory ceiling in GiB; null/omitted means actual node RAM minus 10 GiB |
 | `services.edge.access_log` | `logs/tf-edge-access.jsonl` | TF edge JSONL accounting log |
 | `services.frontend.admin_user` | empty | Admin account used for frontend system-daemon sudo operations |
 

src/thunder_forge/cluster/config.py

@@ -134,6 +134,7 @@ class NodeRuntime:
     log_level: str | None = None
     max_model_memory: str | None = None
     max_process_memory: str | None = None
+    gpu_memory_gb: int | None = None
     max_concurrent_requests: int | None = None
     paged_ssd_cache_dir: str | None = None
     paged_ssd_cache_max_size: str | None = None
@@ -164,6 +165,7 @@ class ServiceConfig:
     olla_local_binary: str = ""
     olla_log_level: str = DEFAULT_OLLA_LOG_LEVEL
     omlx_port: int = DEFAULT_OMLX_PORT
+    omlx_gpu_memory_gb: int | None = None
     edge_access_log: str = DEFAULT_EDGE_ACCESS_LOG
     edge_inspector_enabled: bool = DEFAULT_EDGE_INSPECTOR_ENABLED
     edge_inspector_output_dir: str = DEFAULT_EDGE_INSPECTOR_OUTPUT_DIR
@@ -375,6 +377,14 @@ def _parse_bool(raw: object, *, name: str, default: bool) -> bool:
     raise ValueError(msg)
 
 
+def _parse_gpu_memory_gb(raw: object, *, name: str) -> int | None:
+    if raw is None:
+        return None
+    if type(raw) is not int or raw <= 0:
+        raise ValueError(f"{name} must be a positive integer in GiB or null (RAM minus 10 GiB)")
+    return raw
+
+
 def _parse_optional_positive_int(raw: object, *, name: str) -> int | None:
     if raw is None:
         return None
@@ -461,6 +471,9 @@ def _parse_services(raw: object) -> ServiceConfig:
     olla_version_raw = str(olla_raw.get("version", "")).strip()
     olla_version_pinned = bool(olla_version_raw) or not has_olla_block
     olla_version = olla_version_raw or DEFAULT_OLLA_VERSION
+    omlx_raw = raw.get("omlx", {}) or {}
+    if not isinstance(omlx_raw, dict):
+        raise ValueError("services.omlx must be a mapping")
     return ServiceConfig(
         edge_host=edge_host,
         edge_port=_parse_service_port(raw, "edge", DEFAULT_EDGE_PORT),
@@ -482,6 +495,9 @@ def _parse_services(raw: object) -> ServiceConfig:
             default=DEFAULT_OLLA_LOG_LEVEL,
         ) or DEFAULT_OLLA_LOG_LEVEL,
         omlx_port=_parse_service_port(raw, "omlx", DEFAULT_OMLX_PORT),
+        omlx_gpu_memory_gb=_parse_gpu_memory_gb(
+            omlx_raw.get("gpu_memory_gb"), name="services.omlx.gpu_memory_gb",
+        ),
         edge_access_log=edge_access_log,
         edge_inspector_enabled=_parse_bool(
             edge_inspector_raw.get("enabled"),
@@ -565,6 +581,7 @@ def _parse_node_runtime(
     *,
     node_id: str,
     default_port: int = DEFAULT_OMLX_PORT,
+    default_gpu_memory_gb: int | None = None,
 ) -> NodeRuntime | None:
     if raw is None:
         return None
@@ -594,6 +611,9 @@ def _parse_node_runtime(
             raw.get("max_process_memory"),
             name=f"{path}.max_process_memory",
         ),
+        gpu_memory_gb=_parse_gpu_memory_gb(
+            raw.get("gpu_memory_gb", default_gpu_memory_gb), name=f"{path}.gpu_memory_gb",
+        ),
         max_concurrent_requests=_parse_optional_positive_int(
             raw.get("max_concurrent_requests"),
             name=f"{path}.max_concurrent_requests",
@@ -696,7 +716,10 @@ def parse_cluster_config(raw: dict) -> ClusterConfig:
             admin_user=str(v.get("admin_user", "")).strip(),
             roles=roles,
             fabric_host=_parse_fabric_host(v.get("fabric_host"), node_id=k),
-            runtime=_parse_node_runtime(v.get("runtime"), node_id=k, default_port=services.omlx_port),
+            runtime=_parse_node_runtime(
+                v.get("runtime"), node_id=k, default_port=services.omlx_port,
+                default_gpu_memory_gb=services.omlx_gpu_memory_gb,
+            ),
             models=list(v.get("models", [])),
             platform=v.get("platform"),
             shell=v.get("shell"),

src/thunder_forge/cluster/omlx.py

@@ -793,6 +793,55 @@ def ensure_omlx_tooling(
     return result
 
 
+def _gpu_memory_start_script(node: Node, program_arguments: list[str]) -> str:
+    """Apply the host-wide Metal ceiling before dropping privileges to run oMLX."""
+    if node.runtime is None or not node.user or node.user == "root":
+        raise ValueError("GPU memory setup requires a runtime and a non-root node user")
+    limit = node.runtime.gpu_memory_gb
+    if limit is not None and (type(limit) is not int or limit <= 0):
+        raise ValueError("gpu_memory_gb must be a positive integer in GiB or null")
+    if limit is not None and limit * 1024 > 2**63 - 1:
+        raise ValueError("gpu_memory_gb exceeds the supported system integer range")
+    # GB in tfconfig follows the binary units of hw.memsize and iogpu.wired_limit_mb.
+    target = str(limit * 1024) if limit is not None else "$((total_mb - 10240))"
+    command = shlex.join([
+        "/usr/bin/sudo", "-n", "-H", "-u", node.user, "--", "/usr/bin/env",
+        f"HOME={node.home_dir}",
+        f"PATH={node.home_dir}/.local/bin:/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin",
+        *program_arguments,
+    ])
+    return f"""set -eu
+if [ "$(/usr/bin/uname -s)" != Darwin ] || [ "$(/usr/bin/uname -m)" != arm64 ]; then
+    echo "GPU memory setup requires Apple Silicon macOS" >&2
+    exit 1
+fi
+runtime_uid=$(/usr/bin/id -u {shlex.quote(node.user)})
+if [ "$runtime_uid" -eq 0 ]; then
+    echo "oMLX must run as a non-root user" >&2
+    exit 1
+fi
+total_bytes=$(/usr/sbin/sysctl -n hw.memsize)
+case "$total_bytes" in ''|*[!0-9]*) echo "Invalid hw.memsize" >&2; exit 1;; esac
+total_mb=$((total_bytes / 1048576))
+limit_mb={target}
+if [ "$limit_mb" -le 0 ] || [ "$limit_mb" -ge "$total_mb" ]; then
+    echo "GPU memory limit must be positive and below physical RAM: $limit_mb / $total_mb MiB" >&2
+    exit 1
+fi
+current_mb=$(/usr/sbin/sysctl -n iogpu.wired_limit_mb)
+if [ "$current_mb" != "$limit_mb" ]; then
+    /usr/sbin/sysctl -w "iogpu.wired_limit_mb=$limit_mb"
+fi
+actual_mb=$(/usr/sbin/sysctl -n iogpu.wired_limit_mb)
+if [ "$actual_mb" != "$limit_mb" ]; then
+    echo "GPU memory limit readback mismatch: expected $limit_mb, got $actual_mb MiB" >&2
+    exit 1
+fi
+echo "GPU wired memory limit: $actual_mb MiB (physical RAM: $total_mb MiB)"
+exec {command}
+"""
+
+
 def generate_launchd_plist(node: Node, *, system_daemon: bool = False) -> str:
     """Generate a macOS launchd plist for a node-level oMLX daemon."""
     if node.runtime is None:
@@ -806,6 +855,7 @@ def generate_launchd_plist(node: Node, *, system_daemon: bool = False) -> str:
     program_arguments = _omlx_program_arguments(node)
     environment = {}
     if system_daemon:
+        program_arguments = ["/bin/sh", "-c", _gpu_memory_start_script(node, program_arguments)]
         environment = {
             "HOME": node.home_dir,
             "PATH": f"{node.home_dir}/.local/bin:/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin",
@@ -818,7 +868,7 @@ def generate_launchd_plist(node: Node, *, system_daemon: bool = False) -> str:
         stdout_log=f"{log_dir}/omlx-{node.runtime.port}.stdout.log",
         stderr_log=f"{log_dir}/omlx-{node.runtime.port}.stderr.log",
         environment=environment,
-        user=node.user,
+        user="root" if system_daemon else node.user,
     )
     return generate_service_launchd_plist(spec, system_daemon=system_daemon)
 

tests/test_cli_runtime.py

@@ -212,6 +212,9 @@ def test_runtime_restart_daemon_dry_run_prints_plist_and_sudo_commands(tmp_path:
     config_dir.mkdir()
     (repo / "tfconfig.yaml").write_text(
         dedent("""\
+            services:
+              omlx:
+                gpu_memory_gb: 100
             models: {}
             nodes:
               infer-03:
@@ -238,6 +241,8 @@ def test_runtime_restart_daemon_dry_run_prints_plist_and_sudo_commands(tmp_path:
     assert "plist_path: /Library/LaunchDaemons/com.thunder-forge.omlx-8018.plist" in result.stdout
     assert "staging_plist_path: /Users/shag/.omlx/run/com.thunder-forge.omlx-8018.plist" in result.stdout
     assert "<key>UserName</key>" in result.stdout
+    assert "limit_mb=102400" in result.stdout
+    assert "iogpu.wired_limit_mb" in result.stdout
     assert "sudo -n /usr/bin/install" in result.stdout
     assert "sudo -n /bin/launchctl bootstrap system" in result.stdout
 

tests/test_parse_cluster_config.py

@@ -149,3 +149,30 @@ def test_parse_model_server_args_partial():
     assert sa.prompt_concurrency is None
     assert sa.temp is None
     assert sa.draft_model is None
+
+
+@pytest.mark.parametrize("global_limit,node_limit,expected", [
+    (None, "omitted", None), (118, "omitted", 118), (118, 90, 90), (118, None, None),
+])
+def test_gpu_memory_config_inheritance(global_limit, node_limit, expected):
+    runtime = {"type": "omlx"}
+    if node_limit != "omitted":
+        runtime["gpu_memory_gb"] = node_limit
+    config = parse_cluster_config({
+        "services": {"omlx": {"gpu_memory_gb": global_limit}},
+        "nodes": {"n": {"host": "n.lan", "ram_gb": 128, "runtime": runtime}},
+    })
+    assert config.services.omlx_gpu_memory_gb == global_limit
+    assert config.nodes["n"].runtime.gpu_memory_gb == expected
+
+
+@pytest.mark.parametrize("value", [0, -1, True, 118.5, "118", "auto", "118; id"])
+@pytest.mark.parametrize("scope", ["global", "node"])
+def test_gpu_memory_config_rejects_invalid_values(value, scope):
+    raw = {"nodes": {"n": {"host": "n.lan", "ram_gb": 128, "runtime": {"type": "omlx"}}}}
+    if scope == "global":
+        raw["services"] = {"omlx": {"gpu_memory_gb": value}}
+    else:
+        raw["nodes"]["n"]["runtime"]["gpu_memory_gb"] = value
+    with pytest.raises(ValueError, match="gpu_memory_gb"):
+        parse_cluster_config(raw)

tests/test_runtime_install.py

@@ -2,8 +2,12 @@
 
 from __future__ import annotations
 
+import plistlib
+import shlex
 import subprocess
 
+import pytest
+
 from thunder_forge.cluster.config import RuntimeType
 from thunder_forge.cluster.omlx import (
     Node,
@@ -74,7 +78,10 @@ def test_generate_launchd_plist_system_daemon_runs_as_node_user() -> None:
     plist = generate_launchd_plist(_make_runtime_node(), system_daemon=True)
 
     assert "<key>UserName</key>" in plist
-    assert "<string>shag</string>" in plist
+    parsed = plistlib.loads(plist.encode())
+    assert parsed["UserName"] == "root"
+    assert parsed["ProgramArguments"][:2] == ["/bin/sh", "-c"]
+    assert "exec /usr/bin/sudo -n -H -u shag -- /usr/bin/env" in parsed["ProgramArguments"][2]
     assert "<key>EnvironmentVariables</key>" in plist
     assert "<key>HOME</key>" in plist
     assert "<string>/Users/shag</string>" in plist
@@ -431,3 +438,84 @@ def fake_ssh_run(user, ip, cmd, *, timeout):
     assert len(calls) == 2
     assert "pkill" in calls[0]
     assert "nohup" in calls[1]
+
+
+@pytest.mark.parametrize(
+    "configured,total_gib,current_mb,write_failure,readback_bad,expected_mb,success",
+    [
+        (None, 128, 0, False, False, 120832, True),
+        (None, 64, 0, False, False, 55296, True),
+        (100, 128, 0, False, False, 102400, True),
+        (118, 128, 120832, False, False, 120832, True),
+        (None, 8, 0, False, False, None, False),
+        (128, 128, 0, False, False, None, False),
+        (129, 128, 0, False, False, None, False),
+        (118, 128, 0, True, False, 120832, False),
+        (118, 128, 0, False, True, 120832, False),
+    ],
+)
+def test_gpu_memory_start_script_executes_before_unprivileged_runtime(
+    tmp_path, configured, total_gib, current_mb, write_failure, readback_bad, expected_mb, success,
+):
+    node = _make_runtime_node(home_dir="/Users/model operator")
+    node.runtime.gpu_memory_gb = configured
+    plist = plistlib.loads(generate_launchd_plist(node, system_daemon=True).encode())
+    script = plist["ProgramArguments"][2]
+    state = tmp_path / "limit"
+    state.write_text(str(current_mb))
+    events = tmp_path / "events"
+    sysctl = tmp_path / "sysctl"
+    sysctl.write_text(f"""#!/bin/sh
+set -eu
+case "$*" in
+  '-n hw.memsize') echo {total_gib * 1024**3} ;;
+  '-n iogpu.wired_limit_mb') cat {shlex.quote(str(state))} ;;
+  '-w iogpu.wired_limit_mb='*)
+    echo write >> {shlex.quote(str(events))}
+    {"exit 1" if write_failure else ":"}
+    {":" if readback_bad else 'echo "${2#*=}" > ' + shlex.quote(str(state))}
+    ;;
+  *) exit 2 ;;
+esac
+""")
+    sysctl.chmod(0o700)
+    sudo = tmp_path / "sudo"
+    sudo.write_text(f"""#!/bin/sh
+set -eu
+[ "$1 $2 $3 $4 $5 $6" = '-n -H -u shag -- /usr/bin/env' ]
+[ "$7" = 'HOME=/Users/model operator' ]
+[ "$9" = '/Users/model operator/.local/bin/omlx' ]
+echo start >> {shlex.quote(str(events))}
+""")
+    sudo.chmod(0o700)
+    # Execute only fixture commands: never call real sudo or mutate the host sysctl.
+    script = script.replace("/usr/sbin/sysctl", shlex.quote(str(sysctl)))
+    script = script.replace("/usr/bin/sudo", shlex.quote(str(sudo)))
+    script = script.replace("/usr/bin/id -u shag", "printf 501")
+    script = script.replace("/usr/bin/uname -s", "printf Darwin")
+    script = script.replace("/usr/bin/uname -m", "printf arm64")
+    result = subprocess.run(["/bin/sh", "-c", script], capture_output=True, text=True)
+    assert (result.returncode == 0) == success, result.stderr
+    calls = events.read_text().splitlines() if events.exists() else []
+    assert ("start" in calls) == success
+    if success:
+        assert int(state.read_text()) == expected_mb
+        assert calls == (["start"] if current_mb == expected_mb else ["write", "start"])
+    elif expected_mb is None:
+        assert not calls
+
+
+
+@pytest.mark.parametrize("user", ["root", ""])
+def test_gpu_memory_rejects_invalid_runtime_user(user):
+    node = _make_runtime_node()
+    node.user = user
+    with pytest.raises(ValueError, match="non-root"):
+        generate_launchd_plist(node, system_daemon=True)
+
+
+def test_gpu_memory_rejects_overflow():
+    node = _make_runtime_node()
+    node.runtime.gpu_memory_gb = 2**64
+    with pytest.raises(ValueError, match="integer range"):
+        generate_launchd_plist(node, system_daemon=True)

tfconfig.example.yaml

@@ -26,6 +26,10 @@ services:
       output_dir: logs/inspector
   omlx:
     port: 8018
+    # Host-wide Metal wired-memory ceiling, in GiB (1 GiB = 1024 MiB).
+    # null/omitted: actual node RAM minus 10 GiB, computed at each daemon start.
+    # Applied before oMLX starts by the system LaunchDaemon, including after reboot.
+    gpu_memory_gb: null
 
 # ── Operator workflow defaults ──────────────────────────────────
 # Non-secret command defaults used by `make smoke` and `make sync`.
@@ -169,6 +173,8 @@ nodes:
     admin_user: admin
     runtime:
       type: omlx
+      # Optional override of services.omlx.gpu_memory_gb; null restores RAM minus 10 GiB.
+      # gpu_memory_gb: 118
       trusted_network: true
     models:
       - memory