minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yongji Wu
2026-06-14 00:38:40 -07:00
committed by GitHub
parent b796338271
commit f2d7d67603
3 changed files with 169 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -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
+79 -26
View File
@@ -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."
test/registered/utils/test_numa_utils.py
+89 -36
View File
@@ -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_with_full_affinity(
self, mock_psutil, _mock_isdir, _mock_which, _mock_mempolicy
def test_returns_true_on_numa_system(
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.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()