diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index 6f7deb25a..0cbadd15e 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/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 _handle_numa_bind_failure( - node: int, allowed_cpus, gpu_id: Optional[int] = None -) -> None: - gpu_str = f" for GPU {gpu_id}" if gpu_id is not None else "" - msg = ( - f"NUMA node {node} has no CPU cores allowed by the current affinity " - f"{sorted(allowed_cpus)}, skipping NUMA binding{gpu_str}." +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") ) - logger.warning(msg) + + +def _probe_numactl_args(numactl_args: str) -> tuple[Optional[str], str]: + """Dry-run ``numactl 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 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=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 "" + 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(reason) if envs.SGLANG_CRASH_ON_NUMA_BIND_FAILURE.get(): - raise RuntimeError(msg) + raise RuntimeError(reason) def _can_set_mempolicy() -> bool: diff --git a/test/registered/utils/test_numa_utils.py b/test/registered/utils/test_numa_utils.py index f5f33a3f3..c59c34a44 100644 --- a/test/registered/utils/test_numa_utils.py +++ b/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()