From aa9847d3370fb8e1024116dc8bff8da2926224e3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 08:57:50 +0000 Subject: [PATCH] Studio: cpuset-aware NUMA topology and env-forced CPU offload detection Follow-up review on the CPU-only / NUMA hardening: - Restrict the NUMA topology to the cpuset-allowed nodes (/proc/self/status Mems_allowed_list). numactl --interleave=all only spans those, so under a Docker/systemd cpuset that exposes the host's nodes but allows one, the decision no longer sums RAM the child cannot use and wrap a launch that OOMs. - Honor env-forced CPU offload: _extra_args_forces_cpu_offload now also reads the inherited LLAMA_ARG_N_GPU_LAYERS / LLAMA_ARG_DEVICE the child would apply, so a GPU host running with LLAMA_ARG_N_GPU_LAYERS=0 or LLAMA_ARG_DEVICE=none gets the CPU-only safe defaults. CLI extras still win over the env. Unit tests and the offline simulations updated and passing. --- studio/backend/core/inference/llama_cpp.py | 25 ++++++++--- studio/backend/core/inference/numa.py | 20 +++++++++ .../backend/tests/test_cpu_only_defaults.py | 44 ++++++++++++------- studio/backend/tests/test_numa_interleave.py | 20 +++++++++ 4 files changed, 85 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index a52fbdd91d..bc31b58c16 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1114,13 +1114,16 @@ def _extra_args_n_ubatch( return None -def _extra_args_forces_cpu_offload(extra_args: Optional[Iterable[str]]) -> bool: - """True when extras run the launch fully on CPU despite a visible GPU, so the - CPU-only safe defaults apply: zero GPU layers (-ngl 0 / --n-gpu-layers 0 / - --gpu-layers 0) or no device (--device/-dev none). Each flag's last occurrence - wins (as llama-server parses it); the two controls are independent.""" +def _extra_args_forces_cpu_offload( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True when the launch runs fully on CPU despite a visible GPU, so the CPU-only + safe defaults apply: zero GPU layers (-ngl 0 / --n-gpu-layers 0 / --gpu-layers 0) + or no device (--device/-dev none). CLI extras win (each control's last value); else + the inherited LLAMA_ARG_N_GPU_LAYERS / LLAMA_ARG_DEVICE env the child would honor.""" args = [str(a) for a in extra_args] if extra_args else [] - ngl_zero = device_none = False + ngl_zero: Optional[bool] = None + device_none: Optional[bool] = None for i, raw in enumerate(args): flag, eq, inline = raw.partition("=") value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") @@ -1131,7 +1134,15 @@ def _extra_args_forces_cpu_offload(extra_args: Optional[Iterable[str]]) -> bool: continue elif flag in ("--device", "-dev"): device_none = value.strip().lower() == "none" - return ngl_zero or device_none + _env = os.environ if env is None else env + if ngl_zero is None and _env.get("LLAMA_ARG_N_GPU_LAYERS") is not None: + try: + ngl_zero = int(_env["LLAMA_ARG_N_GPU_LAYERS"]) == 0 + except (TypeError, ValueError): + pass + if device_none is None and _env.get("LLAMA_ARG_DEVICE") is not None: + device_none = _env["LLAMA_ARG_DEVICE"].strip().lower() == "none" + return bool(ngl_zero) or bool(device_none) def _build_ngram_mod_flags( diff --git a/studio/backend/core/inference/numa.py b/studio/backend/core/inference/numa.py index dfdae1e02e..49a88e69e1 100644 --- a/studio/backend/core/inference/numa.py +++ b/studio/backend/core/inference/numa.py @@ -62,6 +62,21 @@ def _parse_online(spec: str) -> list[int]: return out +def _mems_allowed() -> set[int] | None: + """NUMA nodes the process may allocate on (cpuset), from /proc/self/status + Mems_allowed_list. None when unavailable, so callers fall back to the online set. + numactl --interleave=all only spans these, so a cpuset-limited container must not + count host nodes it cannot use.""" + try: + text = Path("/proc/self/status").read_text() + except OSError: + return None + for line in text.splitlines(): + if line.startswith("Mems_allowed_list:"): + return set(_parse_online(line.split(":", 1)[1])) + return None + + def _node_memfree_mib(node: int) -> int | None: """MemFree for one node from /sys/.../nodeN/meminfo, in MiB (kB -> MiB).""" try: @@ -88,8 +103,13 @@ def read_numa_topology() -> NumaTopology: online = (_NODE_ROOT / "online").read_text() except OSError: return NumaTopology() + # Restrict to cpuset-allowed nodes: under a Docker/systemd cpuset the child can only + # allocate on these, so counting other host nodes would overstate the fittable RAM. + allowed = _mems_allowed() free: dict[int, int] = {} for node in _parse_online(online): + if allowed is not None and node not in allowed: + continue mib = _node_memfree_mib(node) if mib is not None: free[node] = mib diff --git a/studio/backend/tests/test_cpu_only_defaults.py b/studio/backend/tests/test_cpu_only_defaults.py index 8c548353d4..9e639d2d04 100644 --- a/studio/backend/tests/test_cpu_only_defaults.py +++ b/studio/backend/tests/test_cpu_only_defaults.py @@ -234,28 +234,38 @@ def test_explicit_user_numa_skips_auto_interleave_prefix(): def test_extra_args_forces_cpu_offload_helper(): - """The CPU-force detector: zero GPU layers or --device none (PR review fixes).""" + """The CPU-force detector: zero GPU layers or --device none, via CLI or the inherited + LLAMA_ARG_* env (PR review fixes).""" from core.inference.llama_cpp import _extra_args_forces_cpu_offload as f - assert f(["-ngl", "0"]) - assert f(["--n-gpu-layers", "0"]) - assert f(["--gpu-layers", "0"]) - assert f(["-ngl=0"]) - assert not f(["-ngl", "99"]) - assert not f([]) - assert not f(None) - assert not f(["--flash-attn", "on"]) + E: dict = {} # explicit empty env so cases ignore the ambient environment + assert f(["-ngl", "0"], env = E) + assert f(["--n-gpu-layers", "0"], env = E) + assert f(["--gpu-layers", "0"], env = E) + assert f(["-ngl=0"], env = E) + assert not f(["-ngl", "99"], env = E) + assert not f([], env = E) + assert not f(None, env = E) + assert not f(["--flash-attn", "on"], env = E) # Each flag's last occurrence wins, matching llama-server's own parsing. - assert f(["-ngl", "99", "-ngl", "0"]) - assert not f(["-ngl", "0", "-ngl", "99"]) + assert f(["-ngl", "99", "-ngl", "0"], env = E) + assert not f(["-ngl", "0", "-ngl", "99"], env = E) # --device/-dev none also forces CPU, independently of -ngl. - assert f(["--device", "none"]) - assert f(["-dev", "none"]) - assert f(["--device=none"]) - assert not f(["--device", "CUDA0"]) + assert f(["--device", "none"], env = E) + assert f(["-dev", "none"], env = E) + assert f(["--device=none"], env = E) + assert not f(["--device", "CUDA0"], env = E) # The two controls are independent: -ngl 0 stays CPU even with a device named. - assert f(["-ngl", "0", "--device", "CUDA0"]) - assert f(["--device", "none", "-ngl", "99"]) + assert f(["-ngl", "0", "--device", "CUDA0"], env = E) + assert f(["--device", "none", "-ngl", "99"], env = E) + # Inherited env forces CPU when the CLI does not set the control. + assert f([], env = {"LLAMA_ARG_N_GPU_LAYERS": "0"}) + assert f([], env = {"LLAMA_ARG_DEVICE": "none"}) + assert not f([], env = {"LLAMA_ARG_N_GPU_LAYERS": "99"}) + assert not f([], env = {"LLAMA_ARG_DEVICE": "CUDA0"}) + # CLI wins over env. + assert not f(["-ngl", "99"], env = {"LLAMA_ARG_N_GPU_LAYERS": "0"}) + assert f(["-ngl", "0"], env = {"LLAMA_ARG_N_GPU_LAYERS": "99"}) def test_cpu_cap_lowers_advertised_ceiling(): diff --git a/studio/backend/tests/test_numa_interleave.py b/studio/backend/tests/test_numa_interleave.py index 8cd4a3770c..7d856327a0 100644 --- a/studio/backend/tests/test_numa_interleave.py +++ b/studio/backend/tests/test_numa_interleave.py @@ -155,6 +155,26 @@ def test_unknown_model_size_does_not_force_interleave(): assert d.interleave is False +def test_topology_restricted_to_cpuset_allowed_nodes(tmp_path, monkeypatch): + """A cpuset that allows only node 0 must drop host node 1 from the topology, so a + container is not told a node's RAM is usable when the child can't allocate there + (PR review fix).""" + import core.inference.numa as m + + (tmp_path / "online").write_text("0-1") + for nid, free_kb in ((0, 100 * 1024), (1, 200 * 1024)): + d = tmp_path / f"node{nid}" + d.mkdir() + (d / "meminfo").write_text(f"Node {nid} MemFree: {free_kb} kB\n") + monkeypatch.setattr(m, "_NODE_ROOT", tmp_path) + + monkeypatch.setattr(m, "_mems_allowed", lambda: {0}) + assert m.read_numa_topology().node_free_mib == {0: 100} + # No cpuset info (None) -> trust the online set, both nodes present. + monkeypatch.setattr(m, "_mems_allowed", lambda: None) + assert m.read_numa_topology().node_free_mib == {0: 100, 1: 200} + + def test_decision_is_frozen_dataclass(): d = InterleaveDecision(True, "x", ("numactl", "--interleave=all")) try: