From 688c23ce489c56cd487c6b3fcb014d1710809b36 Mon Sep 17 00:00:00 2001 From: gnezim Date: Mon, 19 Oct 2026 18:39:18 +0300 Subject: [PATCH 1/2] feat: configure GPU wired memory before starting oMLX --- README.md | 9 +++ src/thunder_forge/cluster/config.py | 25 +++++++- src/thunder_forge/cluster/omlx.py | 52 ++++++++++++++++- tests/test_cli_runtime.py | 5 ++ tests/test_parse_cluster_config.py | 27 +++++++++ tests/test_runtime_install.py | 90 ++++++++++++++++++++++++++++- tfconfig.example.yaml | 6 ++ 7 files changed, 211 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index cbd1834..326700e 100644 --- a/README.md +++ b/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 ` 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..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 `, `make restart `, 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 ` 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 `max_model_memory` or `max_process_memory`; review those separately if they remain 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 | diff --git a/src/thunder_forge/cluster/config.py b/src/thunder_forge/cluster/config.py index c6ae1fc..cfe7787 100644 --- a/src/thunder_forge/cluster/config.py +++ b/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"), diff --git a/src/thunder_forge/cluster/omlx.py b/src/thunder_forge/cluster/omlx.py index f9bac0a..53c01e3 100644 --- a/src/thunder_forge/cluster/omlx.py +++ b/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) diff --git a/tests/test_cli_runtime.py b/tests/test_cli_runtime.py index 2469620..3a31821 100644 --- a/tests/test_cli_runtime.py +++ b/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 "UserName" 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 diff --git a/tests/test_parse_cluster_config.py b/tests/test_parse_cluster_config.py index 6d963e1..af15b23 100644 --- a/tests/test_parse_cluster_config.py +++ b/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) diff --git a/tests/test_runtime_install.py b/tests/test_runtime_install.py index a3348ab..407cc41 100644 --- a/tests/test_runtime_install.py +++ b/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 "UserName" in plist - assert "shag" 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 "EnvironmentVariables" in plist assert "HOME" in plist assert "/Users/shag" 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) diff --git a/tfconfig.example.yaml b/tfconfig.example.yaml index 69ed7f0..a0c9f41 100644 --- a/tfconfig.example.yaml +++ b/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 From 94d3b0f662f09de350a2ee1051337aa3b42f53ed Mon Sep 17 00:00:00 2001 From: gnezim Date: Mon, 19 Oct 2026 19:03:36 +0300 Subject: [PATCH 2/2] docs: clarify oMLX memory guard remains independent --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 326700e..b49f10a 100644 --- a/README.md +++ b/README.md @@ -163,7 +163,7 @@ Set `services.omlx.gpu_memory_gb` in `tfconfig.yaml` for a shared ceiling, or ov `make bootstrap `, `make restart `, 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 ` 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 `max_model_memory` or `max_process_memory`; review those separately if they remain 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). +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.