From f33e9fb3f589719b014cb05f6ca6494c2aa36237 Mon Sep 17 00:00:00 2001 From: gnezim Date: Mon, 19 Oct 2026 19:14:02 +0300 Subject: [PATCH 1/2] fix: use current oMLX memory guard options --- README.md | 4 ++- src/thunder_forge/cluster/config.py | 30 +++++++++++++----- src/thunder_forge/cluster/omlx.py | 8 ++--- tests/test_config.py | 30 +++++++++++++++--- tests/test_omlx.py | 48 ++++++++++++++++------------- tfconfig.example.yaml | 3 ++ 6 files changed, 86 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index cbd1834..7109583 100644 --- a/README.md +++ b/README.md @@ -267,11 +267,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..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..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..host`, `nodes..runtime.port`, and node names. Endpoint names are `-omlx-live`. - Olla model aliases: generated from `models..runtime_model_id` and `nodes..models`. - Olla router defaults: still owned by `thunder_forge.cluster.config.generate_olla_config` rather than a YAML schema. Use `olla smoke --expected-endpoint -omlx-live` or `olla dev-smoke --expected-endpoint -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_` 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 diff --git a/src/thunder_forge/cluster/config.py b/src/thunder_forge/cluster/config.py index c6ae1fc..25eefed 100644 --- a/src/thunder_forge/cluster/config.py +++ b/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 max_concurrent_requests: int | None = None paged_ssd_cache_dir: str | None = None paged_ssd_cache_max_size: str | None = None @@ -581,6 +582,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"), @@ -589,11 +608,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, max_concurrent_requests=_parse_optional_positive_int( raw.get("max_concurrent_requests"), name=f"{path}.max_concurrent_requests", diff --git a/src/thunder_forge/cluster/omlx.py b/src/thunder_forge/cluster/omlx.py index f9bac0a..eaec337 100644 --- a/src/thunder_forge/cluster/omlx.py +++ b/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: diff --git a/tests/test_config.py b/tests/test_config.py index 5cda4b5..cd297da 100644 --- a/tests/test_config.py +++ b/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"): diff --git a/tests/test_omlx.py b/tests/test_omlx.py index d22e2ad..634aa89 100644 --- a/tests/test_omlx.py +++ b/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,28 @@ 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()) + assert plist["ProgramArguments"] == arguments def test_build_omlx_serve_command_includes_explicit_model_dir_only_when_configured() -> None: diff --git a/tfconfig.example.yaml b/tfconfig.example.yaml index 69ed7f0..1187e98 100644 --- a/tfconfig.example.yaml +++ b/tfconfig.example.yaml @@ -169,6 +169,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 trusted_network: true models: - memory From 468881576e600fae2a5f394d46df445a98848ee4 Mon Sep 17 00:00:00 2001 From: Shag Date: Mon, 19 Oct 2026 21:45:45 +0400 Subject: [PATCH 2/2] test: adapt omlx plist assertions to wrapped system daemon script --- tests/test_omlx.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_omlx.py b/tests/test_omlx.py index 634aa89..a80f0c9 100644 --- a/tests/test_omlx.py +++ b/tests/test_omlx.py @@ -56,7 +56,12 @@ def test_memory_settings_reach_process_and_launchd(options: dict, expected: list 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()) - assert plist["ProgramArguments"] == arguments + 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: