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

@@ -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(

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

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():

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: