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}"
-            executable, debug_str = _create_numactl_executable(
-                numactl_args=numactl_args
-            )
-            debug_str += (
-                f", logical_gpu_id={gpu_id}, "
-                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
-                return
+            numactl_args = _numactl_cpu_mem_args(numa_node, gpu_id)
+            if numactl_args is not None:
+                executable, debug_str = _create_numactl_executable(
+                    numactl_args=numactl_args
+                )
+                debug_str += (
+                    f", logical_gpu_id={gpu_id}, "
+                    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
+                    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."