numa: bind within allowed CPUs when affinity is already constrained (#26983)

python/sglang/srt/environ.py

@@ -731,6 +731,7 @@ class Envs:
     # Numa
     SGLANG_NUMA_BIND_V2 = EnvBool(True)
     SGLANG_AUTO_NUMA_BIND = EnvBool(False)
+    SGLANG_CRASH_ON_NUMA_BIND_FAILURE = EnvBool(False)
 
     # Metrics
     SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False)

python/sglang/srt/utils/numa_utils.py

@@ -11,7 +11,6 @@ from contextlib import contextmanager
 from pathlib import Path
 from typing import Optional
 
-import psutil
 import torch
 
 from sglang.srt.environ import envs
@@ -28,18 +27,19 @@ 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_args = f"--cpunodebind={numa_node} --membind={numa_node}"
+            numactl_args = _numactl_cpu_mem_args(numa_node, gpu_id)
-            executable, debug_str = _create_numactl_executable(
+            if numactl_args is not None:
-                numactl_args=numactl_args
+                executable, debug_str = _create_numactl_executable(
-            )
+                    numactl_args=numactl_args
-            debug_str += (
+                )
-                f", logical_gpu_id={gpu_id}, "
+                debug_str += (
-                f"physical_gpu_id={_get_nvml_device_index(gpu_id)}, "
+                    f", logical_gpu_id={gpu_id}, "
-                f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}"
+                    f"physical_gpu_id={_get_nvml_device_index(gpu_id)}, "
-            )
+                    f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}"
-            with _mp_set_executable(executable=executable, debug_str=debug_str):
+                )
-                yield
+                with _mp_set_executable(executable=executable, debug_str=debug_str):
-                return
+                    yield
+                    return
     yield
 
 
@@ -135,9 +135,74 @@ def numa_bind_to_node(node: int):
 
     if libnuma is None or libnuma.numa_available() < 0:
         logger.warning("numa not available on this system, skip bind action")
+        return
+
+    node_cpus = _node_cpus(node)
+    if node_cpus:
+        allowed_cpus = os.sched_getaffinity(0)
+        target_cpus = node_cpus & allowed_cpus
+        if not target_cpus:
+            _handle_numa_bind_failure(node, allowed_cpus)
+            return
+        os.sched_setaffinity(0, target_cpus)
     else:
         libnuma.numa_run_on_node(ctypes.c_int(node))
-        libnuma.numa_set_preferred(ctypes.c_int(node))
+    libnuma.numa_set_preferred(ctypes.c_int(node))
+
+
+class _Bitmask(ctypes.Structure):
+    _fields_ = [("size", ctypes.c_ulong), ("maskp", ctypes.POINTER(ctypes.c_ulong))]
+
+
+def _node_cpus(node: int) -> set:
+    libnuma = get_libnuma()
+    if libnuma is None or libnuma.numa_available() < 0:
+        return set()
+    libnuma.numa_allocate_cpumask.restype = ctypes.POINTER(_Bitmask)
+    libnuma.numa_node_to_cpus.argtypes = [ctypes.c_int, ctypes.POINTER(_Bitmask)]
+    libnuma.numa_node_to_cpus.restype = ctypes.c_int
+    libnuma.numa_bitmask_isbitset.argtypes = [ctypes.POINTER(_Bitmask), ctypes.c_uint]
+    libnuma.numa_bitmask_isbitset.restype = ctypes.c_int
+    libnuma.numa_bitmask_free.argtypes = [ctypes.POINTER(_Bitmask)]
+    mask = libnuma.numa_allocate_cpumask()
+    try:
+        if libnuma.numa_node_to_cpus(node, mask) != 0:
+            return set()
+        return {
+            i
+            for i in range(mask.contents.size)
+            if libnuma.numa_bitmask_isbitset(mask, i)
+        }
+    finally:
+        libnuma.numa_bitmask_free(mask)
+
+
+def _numactl_cpu_mem_args(node: int, gpu_id: int) -> Optional[str]:
+    node_cpus = _node_cpus(node)
+    if not node_cpus:
+        return f"--cpunodebind={node} --membind={node}"
+    allowed_cpus = os.sched_getaffinity(0)
+    target_cpus = node_cpus & allowed_cpus
+    if not target_cpus:
+        _handle_numa_bind_failure(node, allowed_cpus, gpu_id)
+        return None
+    if target_cpus == node_cpus:
+        return f"--cpunodebind={node} --membind={node}"
+    cpu_list = ",".join(str(c) for c in sorted(target_cpus))
+    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}."
+    )
+    logger.warning(msg)
+    if envs.SGLANG_CRASH_ON_NUMA_BIND_FAILURE.get():
+        raise RuntimeError(msg)
 
 
 def _can_set_mempolicy() -> bool:
@@ -166,18 +231,6 @@ def _is_numa_available() -> bool:
     if not os.path.isdir("/sys/devices/system/node/node1"):
         return False
 
-    # Check if affinity is already constrained
-    pid = os.getpid()
-    process = psutil.Process(pid)
-    cpu_affinity = process.cpu_affinity()
-    all_cpus = list(range(psutil.cpu_count()))
-    constrained_affinity = cpu_affinity != all_cpus
-    if constrained_affinity:
-        logger.warning(
-            "NUMA affinity is already constrained for process, skipping NUMA node configuration for GPU. Remove your constraints to allow automatic configuration."
-        )
-        return False
-
     if not shutil.which("numactl") and envs.SGLANG_NUMA_BIND_V2.get():
         logger.debug(
             "numactl command not found, skipping NUMA node configuration for GPU. Install numactl (e.g., apt-get install numactl) to enable automatic NUMA binding."

test/registered/utils/test_numa_utils.py

@@ -1,11 +1,16 @@
 import ctypes
+import os
 import unittest
 from unittest.mock import MagicMock, patch
 
 from sglang.srt.utils.numa_utils import (
+    _handle_numa_bind_failure,
     _is_numa_available,
+    _node_cpus,
+    _numactl_cpu_mem_args,
     _query_numa_node_for_gpu,
     get_numa_node_if_available,
+    numa_bind_to_node,
 )
 from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
 
@@ -26,63 +31,31 @@ class TestIsNumaAvailable(unittest.TestCase):
     def test_returns_false_when_no_numa_nodes(self, _mock_isdir):
         self.assertFalse(_is_numa_available())
 
-    @patch("sglang.srt.utils.numa_utils._is_cuda", True)
-    @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
-    def test_returns_false_when_affinity_constrained(self, mock_psutil, _mock_isdir):
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = [0, 1]
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 128
-
-        self.assertFalse(_is_numa_available())
-
     @patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=True)
     @patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
     @patch("sglang.srt.utils.numa_utils._is_cuda", True)
     @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
+    def test_returns_true_on_numa_system(
-    def test_returns_true_on_numa_system_with_full_affinity(
+        self, _mock_isdir, _mock_which, _mock_mempolicy
-        self, mock_psutil, _mock_isdir, _mock_which, _mock_mempolicy
     ):
-        all_cpus = list(range(128))
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = all_cpus
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 128
-
         self.assertTrue(_is_numa_available())
 
     @patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=False)
     @patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
     @patch("sglang.srt.utils.numa_utils._is_cuda", True)
     @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
     def test_returns_false_when_mempolicy_not_permitted(
-        self, mock_psutil, _mock_isdir, _mock_which, _mock_mempolicy
+        self, _mock_isdir, _mock_which, _mock_mempolicy
     ):
-        all_cpus = list(range(128))
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = all_cpus
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 128
-
         self.assertFalse(_is_numa_available())
 
     @patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=True)
     @patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
     @patch("sglang.srt.utils.numa_utils._is_cuda", True)
     @patch("os.path.isdir", return_value=True)
-    @patch("sglang.srt.utils.numa_utils.psutil")
     def test_isdir_called_with_node1_path(
-        self, mock_psutil, mock_isdir, _mock_which, _mock_mempolicy
+        self, mock_isdir, _mock_which, _mock_mempolicy
     ):
-        all_cpus = list(range(8))
-        mock_process = MagicMock()
-        mock_process.cpu_affinity.return_value = all_cpus
-        mock_psutil.Process.return_value = mock_process
-        mock_psutil.cpu_count.return_value = 8
-
         _is_numa_available()
         mock_isdir.assert_called_with("/sys/devices/system/node/node1")
 
@@ -313,5 +286,85 @@ class TestB200NumaTopology(unittest.TestCase):
         )
 
 
+class TestNumaBindIntersection(unittest.TestCase):
+    """Tests for constraint-aware NUMA binding (node CPUs intersected with the
+    process's allowed CPUs)."""
+
+    @patch("sglang.srt.utils.numa_utils.get_libnuma", return_value=None)
+    def test_node_cpus_no_libnuma_returns_empty(self, _mock_lib):
+        self.assertEqual(_node_cpus(0), set())
+
+    @patch("os.sched_getaffinity", return_value=set(range(72)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_unconstrained_uses_cpunodebind(self, _cpus, _aff):
+        self.assertEqual(_numactl_cpu_mem_args(0, 0), "--cpunodebind=0 --membind=0")
+
+    @patch("os.sched_getaffinity", return_value={0} | set(range(21, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_constrained_uses_physcpubind(self, _cpus, _aff):
+        expected_cpus = ",".join(str(c) for c in [0] + list(range(21, 72)))
+        self.assertEqual(
+            _numactl_cpu_mem_args(0, 0),
+            f"--physcpubind={expected_cpus} --membind=0",
+        )
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"})
+    @patch("os.sched_getaffinity", return_value=set(range(72, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_empty_intersection_returns_none(self, _cpus, _aff):
+        self.assertIsNone(_numactl_cpu_mem_args(0, 0))
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "1"})
+    @patch("os.sched_getaffinity", return_value=set(range(72, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    def test_numactl_args_empty_intersection_crashes_when_enabled(self, _cpus, _aff):
+        with self.assertRaises(RuntimeError):
+            _numactl_cpu_mem_args(0, 0)
+
+    @patch("os.sched_setaffinity")
+    @patch("os.sched_getaffinity", return_value={0} | set(range(21, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    @patch("sglang.srt.utils.numa_utils.get_libnuma")
+    def test_numa_bind_to_node_constrained_sets_intersection(
+        self, mock_libnuma, _cpus, _aff, mock_setaff
+    ):
+        lib = MagicMock()
+        lib.numa_available.return_value = 0
+        mock_libnuma.return_value = lib
+
+        numa_bind_to_node(0)
+
+        mock_setaff.assert_called_once_with(0, {0} | set(range(21, 72)))
+        lib.numa_set_preferred.assert_called_once()
+        lib.numa_run_on_node.assert_not_called()
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"})
+    @patch("os.sched_setaffinity")
+    @patch("os.sched_getaffinity", return_value=set(range(72, 144)))
+    @patch("sglang.srt.utils.numa_utils._node_cpus", return_value=set(range(72)))
+    @patch("sglang.srt.utils.numa_utils.get_libnuma")
+    def test_numa_bind_to_node_empty_intersection_skips(
+        self, mock_libnuma, _cpus, _aff, mock_setaff
+    ):
+        lib = MagicMock()
+        lib.numa_available.return_value = 0
+        mock_libnuma.return_value = lib
+
+        numa_bind_to_node(0)
+
+        mock_setaff.assert_not_called()
+        lib.numa_set_preferred.assert_not_called()
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "1"})
+    def test_handle_failure_raises_when_enabled(self):
+        with self.assertRaises(RuntimeError):
+            _handle_numa_bind_failure(0, {72, 73})
+
+    @patch.dict(os.environ, {"SGLANG_CRASH_ON_NUMA_BIND_FAILURE": "0"})
+    def test_handle_failure_warns_when_disabled(self):
+        with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING"):
+            _handle_numa_bind_failure(0, {72, 73})
+
+
 if __name__ == "__main__":
     unittest.main()