NUMA: probe numactl binding and fall back when --membind is rejected (#28401)

Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>

python/sglang/srt/utils/numa_utils.py

@@ -6,6 +6,7 @@ import multiprocessing
 import os
 import random
 import shutil
+import subprocess
 import time
 from contextlib import contextmanager
 from pathlib import Path
@@ -27,8 +28,34 @@ def configure_subprocess(server_args: ServerArgs, gpu_id: int):
     if envs.SGLANG_NUMA_BIND_V2.get():
         numa_node = get_numa_node_if_available(server_args, gpu_id)
         if numa_node is not None:
+            # _numactl_cpu_mem_args returns None (warn/raise) on empty CPU intersection (#26983).
             numactl_args = _numactl_cpu_mem_args(numa_node, gpu_id)
             if numactl_args is not None:
+                # Verify numactl can actually apply the binding before we exec it
+                # in front of the interpreter; relax the memory policy if not.
+                numactl_args, probe_err = _probe_numactl_args(numactl_args)
+                if numactl_args is None:
+                    # numactl could not apply even a CPU-only binding (e.g.
+                    # set_mempolicy(2)/sched_setaffinity(2) blocked by seccomp,
+                    # which the read-only get_mempolicy(2) probe in
+                    # _can_set_mempolicy cannot detect). Reuse #26983's failure
+                    # semantics (warn-and-continue, or raise when
+                    # SGLANG_CRASH_ON_NUMA_BIND_FAILURE) with an explicit reason
+                    # carrying the captured stderr: the CPU intersection already
+                    # succeeded here, so the default "no CPU cores allowed"
+                    # message would mislead operators toward the wrong cause.
+                    probe_suffix = f": {probe_err}" if probe_err else ""
+                    _handle_numa_bind_failure(
+                        numa_node,
+                        reason=(
+                            f"numactl could not apply NUMA binding for node "
+                            f"{numa_node} (e.g. set_mempolicy/sched_setaffinity "
+                            f"blocked by seccomp, or cpuset rejects the policy)"
+                            f"{probe_suffix}; skipping NUMA binding for GPU {gpu_id}."
+                        ),
+                    )
+                    yield
+                    return
                 executable, debug_str = _create_numactl_executable(
                     numactl_args=numactl_args
                 )
@@ -192,17 +219,120 @@ def _numactl_cpu_mem_args(node: int, gpu_id: int) -> Optional[str]:
     return f"--physcpubind={cpu_list} --membind={node}"
 
 
+def _strip_memory_args(numactl_args: str) -> str:
+    """Return ``numactl_args`` with the ``--membind`` segment removed, keeping
+    only the CPU binding (``--cpunodebind`` / ``--physcpubind``)."""
+    return " ".join(
+        token for token in numactl_args.split() if not token.startswith("--membind")
+    )
+
+
+def _probe_numactl_args(numactl_args: str) -> tuple[Optional[str], str]:
+    """Dry-run ``numactl <args> true`` and fall back to a weaker binding when the
+    kernel rejects the strongest one.
+
+    ``configure_subprocess`` applies NUMA binding by exec-ing ``numactl`` in front
+    of the Python interpreter (see ``_create_numactl_executable``), so a binding
+    that ``numactl`` refuses kills the worker before Python starts, with no
+    traceback. ``_can_set_mempolicy`` only probes ``get_mempolicy(2)`` (read),
+    which does not catch ``set_mempolicy(2)`` being denied (e.g. by a seccomp
+    profile) or a ``--membind`` that the cpuset rejects with ``EINVAL``.
+
+    To avoid that silent crash we probe the requested args and progressively relax
+    the *memory* policy while keeping the CPU binding intact::
+
+        --membind=N  ->  --preferred=N  ->  drop the memory segment
+
+    Returns ``(args, last_stderr)``: ``args`` is the strongest binding that
+    actually runs, or ``None`` if even CPU-only fails (or ``numactl`` is missing /
+    errors out); ``last_stderr`` is the rejection reason numactl printed for the
+    strongest binding that was rejected (empty on success), so the caller can
+    surface it on the total-failure path.
+    """
+
+    def _probe(args: str):
+        """Run ``numactl <args> true``; return ``(succeeded, stderr_text)``."""
+        try:
+            proc = subprocess.run(
+                ["numactl", *args.split(), "true"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.PIPE,
+                timeout=10,
+            )
+            stderr = proc.stderr.decode("utf-8", errors="replace").strip()
+            if proc.returncode != 0:
+                logger.debug(f"numactl probe for {args!r} rejected: {stderr!r}")
+            return proc.returncode == 0, stderr
+        except Exception as e:
+            # Missing numactl, timeout, etc. Treat as "this binding does not work".
+            logger.debug(f"numactl probe for {args!r} failed: {e}")
+            return False, str(e)
+
+    def _suffix(err: str) -> str:
+        return f": {err}" if err else ""
+
+    # 1. Strongest binding: exactly what was requested.
+    ok, last_err = _probe(numactl_args)
+    if ok:
+        return numactl_args, ""
+
+    # 2. Relax a hard --membind=N to a soft --preferred=N. The memory segment here
+    #    is always a single node, which maps cleanly onto --preferred (single-node
+    #    only). MPOL_PREFERRED is a hint and can succeed where MPOL_BIND is denied.
+    if "--membind=" in numactl_args:
+        preferred_args = numactl_args.replace("--membind=", "--preferred=")
+        ok, _ = _probe(preferred_args)
+        if ok:
+            logger.warning(
+                f"numactl rejected hard memory binding ({numactl_args!r})"
+                f"{_suffix(last_err)}; falling back to soft preferred policy "
+                f"({preferred_args!r})."
+            )
+            return preferred_args, ""
+
+    # 3. Drop the memory segment entirely, keep only the CPU binding.
+    cpu_only_args = _strip_memory_args(numactl_args)
+    if cpu_only_args and cpu_only_args != numactl_args:
+        ok, cpu_err = _probe(cpu_only_args)
+        if ok:
+            logger.warning(
+                f"numactl rejected memory binding ({numactl_args!r})"
+                f"{_suffix(last_err)}; falling back to CPU-only binding "
+                f"({cpu_only_args!r})."
+            )
+            return cpu_only_args, ""
+        last_err = cpu_err
+
+    # 4. Nothing worked.
+    return None, last_err
+
+
 def _handle_numa_bind_failure(
-    node: int, allowed_cpus, gpu_id: Optional[int] = None
+    node: int,
+    allowed_cpus=None,
+    gpu_id: Optional[int] = None,
+    *,
+    reason: Optional[str] = None,
 ) -> None:
+    """Emit the NUMA-bind failure warning, or raise it when
+    ``SGLANG_CRASH_ON_NUMA_BIND_FAILURE`` is set.
+
+    Two call modes:
+      * ``reason is None`` (default): the failure is an empty CPU intersection,
+        so the message reports ``allowed_cpus`` (which must be provided).
+      * ``reason`` provided: the failure is something else (e.g. numactl rejected
+        the binding at runtime); the caller supplies the exact message and
+        ``allowed_cpus`` / ``gpu_id`` are not needed.
+    """
+    if reason is None:
         gpu_str = f" for GPU {gpu_id}" if gpu_id is not None else ""
-    msg = (
+        reason = (
             f"NUMA node {node} has no CPU cores allowed by the current affinity "
             f"{sorted(allowed_cpus)}, skipping NUMA binding{gpu_str}."
         )
-    logger.warning(msg)
+    logger.warning(reason)
     if envs.SGLANG_CRASH_ON_NUMA_BIND_FAILURE.get():
-        raise RuntimeError(msg)
+        raise RuntimeError(reason)
 
 
 def _can_set_mempolicy() -> bool:

test/registered/utils/test_numa_utils.py

@@ -1,6 +1,7 @@
 import ctypes
 import os
 import unittest
+from contextlib import ExitStack
 from unittest.mock import MagicMock, patch
 
 from sglang.srt.utils.numa_utils import (
@@ -8,7 +9,10 @@ from sglang.srt.utils.numa_utils import (
     _is_numa_available,
     _node_cpus,
     _numactl_cpu_mem_args,
+    _probe_numactl_args,
     _query_numa_node_for_gpu,
+    _strip_memory_args,
+    configure_subprocess,
     get_numa_node_if_available,
     numa_bind_to_node,
 )
@@ -366,5 +370,201 @@ class TestNumaBindIntersection(unittest.TestCase):
             _handle_numa_bind_failure(0, {72, 73})
 
 
+def _run_result(returncode, stderr=b""):
+    """Build a fake subprocess.CompletedProcess-like object with a returncode
+    and captured stderr (bytes, as subprocess.run(..., stderr=PIPE) returns)."""
+    result = MagicMock()
+    result.returncode = returncode
+    result.stderr = stderr
+    return result
+
+
+class TestProbeNumactlArgs(unittest.TestCase):
+    """Tests for _probe_numactl_args: dry-run numactl and relax the memory policy
+    (--membind -> --preferred -> CPU-only) when the kernel rejects the binding.
+
+    subprocess.run is mocked and orchestrated by returncode; no real numactl or
+    GPU is required. Returns ``(args, last_stderr)``."""
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_membind_probe_succeeds_returns_original(self, mock_run):
+        # The requested binding works on the first probe.
+        mock_run.side_effect = [_run_result(0)]
+        args = "--cpunodebind=0 --membind=0"
+        self.assertEqual(_probe_numactl_args(args), (args, ""))
+        self.assertEqual(mock_run.call_count, 1)
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_membind_fails_preferred_succeeds(self, mock_run):
+        # --membind rejected, --preferred accepted.
+        mock_run.side_effect = [_run_result(1), _run_result(0)]
+        with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
+            result = _probe_numactl_args("--cpunodebind=0 --membind=0")
+        self.assertEqual(result, ("--cpunodebind=0 --preferred=0", ""))
+        self.assertTrue(any("preferred" in msg for msg in cm.output))
+        # Second probe must have used the --preferred form.
+        second_call_argv = mock_run.call_args_list[1].args[0]
+        self.assertIn("--preferred=0", second_call_argv)
+        self.assertNotIn("--membind=0", second_call_argv)
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_membind_and_preferred_fail_cpu_only_succeeds(self, mock_run):
+        # --membind and --preferred rejected, CPU-only accepted.
+        mock_run.side_effect = [_run_result(1), _run_result(1), _run_result(0)]
+        with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
+            result = _probe_numactl_args("--physcpubind=0,21,22 --membind=0")
+        self.assertEqual(result, ("--physcpubind=0,21,22", ""))
+        self.assertTrue(any("CPU-only" in msg for msg in cm.output))
+        third_call_argv = mock_run.call_args_list[2].args[0]
+        self.assertNotIn("--membind=0", third_call_argv)
+        self.assertNotIn("--preferred=0", third_call_argv)
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_all_probes_fail_returns_none_with_last_stderr(self, mock_run):
+        # Every binding, down to CPU-only, is rejected; the returned stderr is the
+        # CPU-only (last / strongest-attempted) rejection reason.
+        mock_run.side_effect = [
+            _run_result(1, stderr=b"numactl: setting membind: Invalid argument"),
+            _run_result(1, stderr=b"numactl: setting preferred: Invalid argument"),
+            _run_result(1, stderr=b"numactl: cpunodebind: Operation not permitted"),
+        ]
+        args, err = _probe_numactl_args("--cpunodebind=0 --membind=0")
+        self.assertIsNone(args)
+        self.assertIn("cpunodebind", err)
+        self.assertEqual(mock_run.call_count, 3)
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_cpu_only_input_failure_returns_none(self, mock_run):
+        # No --membind in the input: the requested args are already CPU-only, so
+        # step 1 is the only probe and on failure we skip --preferred / strip.
+        mock_run.side_effect = [
+            _run_result(1, stderr=b"numactl: cpunodebind: Operation not permitted")
+        ]
+        args, err = _probe_numactl_args("--cpunodebind=0")
+        self.assertIsNone(args)
+        self.assertIn("cpunodebind", err)
+        self.assertEqual(mock_run.call_count, 1)
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_numactl_missing_returns_none(self, mock_run):
+        # numactl not installed / raises: probe must not propagate, returns None.
+        mock_run.side_effect = FileNotFoundError("numactl")
+        args, _err = _probe_numactl_args("--cpunodebind=0 --membind=0")
+        self.assertIsNone(args)
+
+    @patch("sglang.srt.utils.numa_utils.subprocess.run")
+    def test_rejection_stderr_surfaces_in_fallback_warning(self, mock_run):
+        # numactl prints the precise rejection reason to stderr (e.g.
+        # "setting membind: Invalid argument"); the fallback warning must
+        # surface it so operators can tell seccomp vs cpuset apart.
+        mock_run.side_effect = [
+            _run_result(1, stderr=b"numactl: setting membind: Invalid argument"),
+            _run_result(0),
+        ]
+        with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
+            result = _probe_numactl_args("--cpunodebind=0 --membind=0")
+        self.assertEqual(result, ("--cpunodebind=0 --preferred=0", ""))
+        self.assertTrue(
+            any("Invalid argument" in msg for msg in cm.output),
+            f"expected numactl stderr in warning, got {cm.output}",
+        )
+
+
+class TestStripMemoryArgs(unittest.TestCase):
+    """Direct tests for _strip_memory_args: drop --membind, keep CPU binding."""
+
+    def test_strips_membind_keeps_cpu(self):
+        self.assertEqual(
+            _strip_memory_args("--cpunodebind=0 --membind=0"),
+            "--cpunodebind=0",
+        )
+        self.assertEqual(
+            _strip_memory_args("--physcpubind=0,21,22 --membind=0"),
+            "--physcpubind=0,21,22",
+        )
+
+    def test_no_membind_returns_unchanged(self):
+        self.assertEqual(_strip_memory_args("--cpunodebind=0"), "--cpunodebind=0")
+
+
+class TestConfigureSubprocessProbeFailure(unittest.TestCase):
+    """Tests the wiring in configure_subprocess when _probe_numactl_args gives up
+    (returns None): the worker must start unbound (warn-and-yield) by default, or
+    raise before yielding when SGLANG_CRASH_ON_NUMA_BIND_FAILURE=1.
+
+    get_numa_node_if_available / _numactl_cpu_mem_args / _probe_numactl_args are
+    mocked to drive the probe-failure branch directly; _create_numactl_executable
+    and _mp_set_executable are mocked to assert the failure path never installs a
+    numactl executable. No real numactl or GPU is required."""
+
+    def _common_patches(self):
+        return [
+            patch(
+                "sglang.srt.utils.numa_utils.get_numa_node_if_available",
+                return_value=0,
+            ),
+            patch(
+                "sglang.srt.utils.numa_utils._numactl_cpu_mem_args",
+                return_value="--cpunodebind=0 --membind=0",
+            ),
+            patch(
+                "sglang.srt.utils.numa_utils._probe_numactl_args",
+                return_value=(
+                    None,
+                    "numactl: setting membind: Invalid argument",
+                ),
+            ),
+            patch("sglang.srt.utils.numa_utils._create_numactl_executable"),
+            patch("sglang.srt.utils.numa_utils._mp_set_executable"),
+        ]
+
+    @patch.dict(
+        os.environ,
+        {"SGLANG_NUMA_BIND_V2": "1", "SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"},
+    )
+    def test_probe_none_warns_and_yields_unbound(self):
+        with ExitStack() as stack:
+            mocks = [stack.enter_context(p) for p in self._common_patches()]
+            _mock_get, _mock_args, _mock_probe, mock_create, mock_mp = mocks
+            server_args = MagicMock()
+            with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm:
+                with configure_subprocess(server_args, 0):
+                    pass  # worker would start unbound here
+            # The probe-failure path reuses #26983's failure helper (warn) and
+            # must NOT install a numactl executable. The captured numactl stderr
+            # is threaded into the warning so operators can see the rejection cause.
+            self.assertTrue(
+                any("could not apply NUMA binding" in msg for msg in cm.output),
+                f"expected probe-failure warning, got {cm.output}",
+            )
+            self.assertTrue(
+                any("Invalid argument" in msg for msg in cm.output),
+                f"expected numactl stderr in warning, got {cm.output}",
+            )
+            mock_create.assert_not_called()
+            mock_mp.assert_not_called()
+
+    @patch.dict(
+        os.environ,
+        {"SGLANG_NUMA_BIND_V2": "1", "SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "1"},
+    )
+    def test_probe_none_raises_before_yield_when_crash_enabled(self):
+        with ExitStack() as stack:
+            mocks = [stack.enter_context(p) for p in self._common_patches()]
+            _mock_get, _mock_args, _mock_probe, mock_create, mock_mp = mocks
+            server_args = MagicMock()
+            with self.assertRaises(RuntimeError) as cm:
+                with configure_subprocess(server_args, 0):
+                    self.fail(
+                        "contextmanager must not yield when crash-on-failure is set"
+                    )
+            # The RuntimeError carries the captured stderr so crash logs show the
+            # rejection cause, not just the failure category.
+            self.assertIn("could not apply NUMA binding", str(cm.exception))
+            self.assertIn("Invalid argument", str(cm.exception))
+            mock_create.assert_not_called()
+            mock_mp.assert_not_called()
+
+
 if __name__ == "__main__":
     unittest.main()