diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 79d2a922a..209c56fd0 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index 9015f2fbe..6f7deb25a 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/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." diff --git a/test/registered/utils/test_numa_utils.py b/test/registered/utils/test_numa_utils.py index 2653dd81f..f5f33a3f3 100644 --- a/test/registered/utils/test_numa_utils.py +++ b/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_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()