fix: use current oMLX memory guard options

README.md

@@ -276,11 +276,13 @@ Run `uv run thunder-forge config lint` before generating runtime/router config.
 ### Parameter Sources
 
 - Cache artifact root: `.env` key `TF_CACHE_OMLX_MODELS_DIR`, default `~/.omlx/models`. Artifact `status`, `download`, and `sync` use this path on the cache execution host (local cache role or remotely dispatched cache role).
-- Node oMLX process args: `tfconfig.yaml` under `nodes.<node>.runtime`. `type` and `port` are required; optional keys map directly to `omlx serve` flags: `model_dir`, `bind_host`, `base_path`, `log_level`, `max_model_memory`, `max_process_memory`, `max_concurrent_requests`, `paged_ssd_cache_dir`, `paged_ssd_cache_max_size`, `hot_cache_max_size`, `no_cache`, `mcp_config`, and `hf_endpoint`.
+- Node oMLX process args: `tfconfig.yaml` under `nodes.<node>.runtime`. `type` and `port` are required; optional keys map directly to `omlx serve` flags: `model_dir`, `bind_host`, `base_path`, `log_level`, `memory_guard`, `memory_guard_gb`, `max_concurrent_requests`, `paged_ssd_cache_dir`, `paged_ssd_cache_max_size`, `hot_cache_max_size`, `no_cache`, `mcp_config`, and `hf_endpoint`.
 - Olla generated endpoints: `generate-olla-config` reads `nodes.<node>.host`, `nodes.<node>.runtime.port`, and node names. Endpoint names are `<node>-omlx-live`.
 - Olla model aliases: generated from `models.<alias>.runtime_model_id` and `nodes.<node>.models`.
 - Olla router defaults: still owned by `thunder_forge.cluster.config.generate_olla_config` rather than a YAML schema. Use `olla smoke --expected-endpoint <node>-omlx-live` or `olla dev-smoke --expected-endpoint <node>-omlx-live` when you want smoke tests to pin a specific generated endpoint.
 
+oMLX memory settings follow the [current CLI contract](https://github.com/jundot/omlx/blob/7a5a3d75ce7463300bdfa5fedccd34348f60ff46/omlx/cli.py#L1246), verified against `v0.7.1rc1`. `memory_guard` accepts `safe`, `balanced`, `aggressive`, or the quoted YAML string `"off"`. `memory_guard_gb` accepts a finite positive number in oMLX's GB units, including fractions; it selects a custom ceiling and enables the guard. A custom ceiling takes precedence over a tier and cannot be combined with `"off"`. Omitted or null settings add no CLI overrides, preserving oMLX's saved settings/defaults. The obsolete `max_model_memory` and `max_process_memory` keys now fail config validation: remove them and choose a guard tier or custom ceiling explicitly. Their meanings are different, so TF does not convert them automatically. This configures oMLX's memory guard; it does not change the macOS GPU wired-memory ceiling.
+
 With no `.env` and no TF-specific environment variables, Thunder Forge uses these defaults: cache artifacts under `~/.omlx/models`; omitted node users from `GATEWAY_SSH_USER`, then `$USER`, then `unknown`; and no TF edge clients because no `TF_USER_<CLIENT>` entries are set. Commands that need edge auth, such as `edge smoke`, fail until the requested client has a matching key such as `TF_USER_HINDSIGHT`.
 
 ## Topology and Rollout

src/thunder_forge/cluster/config.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import math
 import os
 import subprocess
 from dataclasses import dataclass, field
@@ -132,8 +133,8 @@ class NodeRuntime:
     bind_host: str = "0.0.0.0"
     base_path: str | None = None
     log_level: str | None = None
-    max_model_memory: str | None = None
-    max_process_memory: str | None = None
+    memory_guard: str | None = None
+    memory_guard_gb: float | None = None
     gpu_memory_gb: int | None = None
     max_concurrent_requests: int | None = None
     paged_ssd_cache_dir: str | None = None
@@ -598,6 +599,24 @@ def _parse_node_runtime(
     except ValueError as exc:
         msg = f"{path}.type must be one of: {', '.join(item.value for item in RuntimeType)}"
         raise ValueError(msg) from exc
+    for obsolete in ("max_model_memory", "max_process_memory"):
+        if obsolete in raw:
+            raise ValueError(
+                f"{path}.{obsolete} is no longer supported by oMLX; remove it and configure "
+                "memory_guard or memory_guard_gb explicitly (the old limits are not equivalent)"
+            )
+    memory_guard = _parse_optional_non_empty_string(raw.get("memory_guard"), name=f"{path}.memory_guard")
+    if memory_guard not in {None, "off", "safe", "balanced", "aggressive"}:
+        raise ValueError(f"{path}.memory_guard must be one of: off, safe, balanced, aggressive")
+    memory_guard_gb = None
+    if raw.get("memory_guard_gb") is not None:
+        if isinstance(raw["memory_guard_gb"], bool):
+            raise ValueError(f"{path}.memory_guard_gb must be a finite positive number")
+        memory_guard_gb = _parse_positive_float(raw["memory_guard_gb"], name=f"{path}.memory_guard_gb", default=0)
+        if not math.isfinite(memory_guard_gb):
+            raise ValueError(f"{path}.memory_guard_gb must be a finite positive number")
+        if memory_guard == "off":
+            raise ValueError(f"{path}.memory_guard off cannot be combined with memory_guard_gb")
     return NodeRuntime(
         type=runtime_type,
         port=parse_port(raw.get("port", default_port), name=f"{path}.port"),
@@ -606,11 +625,8 @@ def _parse_node_runtime(
         or "0.0.0.0",
         base_path=_parse_optional_non_empty_string(raw.get("base_path"), name=f"{path}.base_path"),
         log_level=_parse_optional_non_empty_string(raw.get("log_level"), name=f"{path}.log_level"),
-        max_model_memory=_parse_optional_non_empty_string(raw.get("max_model_memory"), name=f"{path}.max_model_memory"),
-        max_process_memory=_parse_optional_non_empty_string(
-            raw.get("max_process_memory"),
-            name=f"{path}.max_process_memory",
-        ),
+        memory_guard=memory_guard,
+        memory_guard_gb=memory_guard_gb,
         gpu_memory_gb=_parse_gpu_memory_gb(
             raw.get("gpu_memory_gb", default_gpu_memory_gb), name=f"{path}.gpu_memory_gb",
         ),

src/thunder_forge/cluster/omlx.py

@@ -160,10 +160,10 @@ def _omlx_program_arguments(node: Node) -> list[str]:
         args.extend(["--base-path", runtime.base_path])
     if runtime.log_level is not None:
         args.extend(["--log-level", runtime.log_level])
-    if runtime.max_model_memory is not None:
-        args.extend(["--max-model-memory", runtime.max_model_memory])
-    if runtime.max_process_memory is not None:
-        args.extend(["--max-process-memory", runtime.max_process_memory])
+    if runtime.memory_guard is not None:
+        args.extend(["--memory-guard", runtime.memory_guard])
+    if runtime.memory_guard_gb is not None:
+        args.extend(["--memory-guard-gb", str(runtime.memory_guard_gb)])
     if runtime.max_concurrent_requests is not None:
         args.extend(["--max-concurrent-requests", str(runtime.max_concurrent_requests)])
     if runtime.paged_ssd_cache_dir is not None:

tests/test_config.py

@@ -376,8 +376,8 @@ def test_parse_node_runtime_options_for_omlx_serve(tmp_path: Path) -> None:
               bind_host: 127.0.0.1
               base_path: /Users/shag/.omlx-tf
               log_level: warning
-              max_model_memory: 90GB
-              max_process_memory: auto
+              memory_guard: balanced
+              memory_guard_gb: 90.5
               max_concurrent_requests: 4
               paged_ssd_cache_dir: /Users/shag/.omlx/cache
               paged_ssd_cache_max_size: 50GB
@@ -397,8 +397,8 @@ def test_parse_node_runtime_options_for_omlx_serve(tmp_path: Path) -> None:
     assert runtime.bind_host == "127.0.0.1"
     assert runtime.base_path == "/Users/shag/.omlx-tf"
     assert runtime.log_level == "warning"
-    assert runtime.max_model_memory == "90GB"
-    assert runtime.max_process_memory == "auto"
+    assert runtime.memory_guard == "balanced"
+    assert runtime.memory_guard_gb == 90.5
     assert runtime.max_concurrent_requests == 4
     assert runtime.paged_ssd_cache_dir == "/Users/shag/.omlx/cache"
     assert runtime.paged_ssd_cache_max_size == "50GB"
@@ -426,6 +426,28 @@ def test_parse_node_runtime_requires_type() -> None:
         )
 
 
+@pytest.mark.parametrize("options", [
+    {"max_model_memory": "90GB"},
+    {"max_process_memory": "auto"},
+    {"max_model_memory": None},
+    {"memory_guard": "custom"},
+    {"memory_guard": False},
+    {"memory_guard_gb": 0},
+    {"memory_guard_gb": -1},
+    {"memory_guard_gb": True},
+    {"memory_guard_gb": float("inf")},
+    {"memory_guard_gb": float("nan")},
+    {"memory_guard_gb": "90GB"},
+    {"memory_guard": "off", "memory_guard_gb": 90},
+])
+def test_parse_node_runtime_rejects_invalid_memory_settings(options: dict) -> None:
+    with pytest.raises(ValueError, match="nodes.infer-03.runtime.*memory_guard"):
+        parse_cluster_config({"nodes": {"infer-03": {
+            "host": "infer-03.lan", "ram_gb": 128,
+            "runtime": {"type": "omlx", **options},
+        }}})
+
+
 @pytest.mark.parametrize("field", ["no_cache", "trusted_network"])
 def test_parse_node_runtime_rejects_string_booleans(field: str) -> None:
     with pytest.raises(ValueError, match=rf"nodes\.infer-03\.runtime\.{field} must be boolean true/false"):

tests/test_omlx.py

@@ -1,15 +1,19 @@
 """Tests for oMLX node-level runtime helpers."""
 
+import plistlib
+import shlex
 from types import SimpleNamespace
 
 import httpx
+import pytest
 
-from thunder_forge.cluster.config import Node, NodeRuntime
+from thunder_forge.cluster.config import Node, NodeRuntime, parse_cluster_config
 from thunder_forge.cluster.omlx import (
     build_omlx_auth_settings_command,
     build_omlx_serve_command,
     check_omlx_health,
     ensure_omlx_auth_settings,
+    generate_launchd_plist,
     resolve_latest_omlx_release,
     smoke_omlx_chat,
 )
@@ -31,26 +35,33 @@ def test_build_omlx_serve_command_omits_default_model_dir() -> None:
     assert "--model-dir" not in command
 
 
-def test_build_omlx_serve_command_includes_memory_caps_when_configured() -> None:
-    node = Node(
-        host="infer-03.lan",
-        ram_gb=128,
-        user="shag",
-        roles=["inference"],
-        runtime=NodeRuntime(
-            type="omlx",
-            port=8018,
-            max_model_memory="90GB",
-            max_process_memory="auto",
-        ),
-        home_dir="/Users/shag",
-    )
-
-    command = build_omlx_serve_command(node)
-
-    assert command == (
-        "/Users/shag/.local/bin/omlx serve --host 0.0.0.0 --port 8018 --max-model-memory 90GB --max-process-memory auto"
-    )
+@pytest.mark.parametrize("options,expected", [
+    ({}, []),
+    ({"memory_guard": "off"}, ["--memory-guard", "off"]),
+    ({"memory_guard": "safe"}, ["--memory-guard", "safe"]),
+    ({"memory_guard": "balanced"}, ["--memory-guard", "balanced"]),
+    ({"memory_guard": "aggressive"}, ["--memory-guard", "aggressive"]),
+    ({"memory_guard_gb": 90.5}, ["--memory-guard-gb", "90.5"]),
+    ({"memory_guard": "safe", "memory_guard_gb": 90},
+     ["--memory-guard", "safe", "--memory-guard-gb", "90.0"]),
+])
+def test_memory_settings_reach_process_and_launchd(options: dict, expected: list[str]) -> None:
+    config = parse_cluster_config({"nodes": {"infer-03": {
+        "host": "infer-03.lan", "ram_gb": 128, "user": "shag",
+        "runtime": {"type": "omlx", **options},
+    }}})
+    node = config.nodes["infer-03"]
+    node.home_dir = "/Users/shag"
+    arguments = ["/Users/shag/.local/bin/omlx", "serve", "--host", "0.0.0.0", "--port", "8018", *expected]
+    assert shlex.split(build_omlx_serve_command(node)) == arguments
+    for system_daemon in (False, True):
+        plist = plistlib.loads(generate_launchd_plist(node, system_daemon=system_daemon).encode())
+        if system_daemon:
+            # PR #5 wraps system-daemon arguments in a root sh -c GPU-memory setup script.
+            assert plist["ProgramArguments"][:2] == ["/bin/sh", "-c"]
+            assert shlex.join(arguments) in plist["ProgramArguments"][2]
+        else:
+            assert plist["ProgramArguments"] == arguments
 
 
 def test_build_omlx_serve_command_includes_explicit_model_dir_only_when_configured() -> None:

tfconfig.example.yaml

@@ -173,6 +173,9 @@ nodes:
     admin_user: admin
     runtime:
       type: omlx
+      # Optional oMLX memory guard; omit to preserve saved settings/defaults.
+      # memory_guard: balanced  # safe, balanced, aggressive, or quoted "off"
+      # memory_guard_gb: 100.5  # custom GB ceiling; enables the guard, overrides the tier
       # Optional override of services.omlx.gpu_memory_gb; null restores RAM minus 10 GiB.
       # gpu_memory_gb: 118
       trusted_network: true