diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 539dff12a..60043fc49 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -207,10 +207,8 @@ from sglang.srt.utils import ( get_available_gpu_memory, get_bool_env_var, get_int_env_var, - get_numa_node, is_mps, kill_itself_when_parent_died, - numa_bind_to_node, point_to_point_pyobj, require_mlp_sync, set_gpu_proc_affinity, @@ -224,6 +222,7 @@ from sglang.srt.utils.hf_transformers_utils import ( get_tokenizer_from_processor, ) from sglang.srt.utils.network import get_zmq_socket +from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node from sglang.srt.utils.tensor_bridge import use_mlx from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.utils import TypeBasedDispatcher, get_exception_traceback @@ -3542,12 +3541,7 @@ def run_scheduler_process( set_gpu_proc_affinity( server_args.pp_size, server_args.tp_size, server_args.nnodes, gpu_id ) - numa_node = None - if (numa_nodes := server_args.numa_node) is not None: - numa_node = numa_nodes[gpu_id] - elif envs.SGLANG_AUTO_NUMA_BIND.get(): - numa_node = get_numa_node(gpu_id) - logger.info(f"auto get NUMA node {numa_node} for GPU {gpu_id}") + numa_node = get_numa_node_if_available(server_args, gpu_id) if numa_node is not None and not envs.SGLANG_NUMA_BIND_V2.get(): numa_bind_to_node(numa_node) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index f96c5ac8b..6130d9502 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -47,7 +47,6 @@ from sglang.srt.mem_cache.radix_cache import ( split_node_hash_value, ) from sglang.srt.observability.metrics_collector import StorageMetricsCollector -from sglang.srt.utils import bind_to_closest_numa_node_cuda if TYPE_CHECKING: from sglang.srt.mem_cache.cache_init_params import CacheInitParams @@ -104,9 +103,6 @@ class HiMambaRadixCache(MambaRadixCache): "switching to page first direct layout" ) - if not server_args.disable_hicache_numa_detect: - bind_to_closest_numa_node_cuda() - self.page_size = params.page_size self.hybrid_kv_cache = params.token_to_kv_pool_allocator.get_kvcache() if not isinstance(self.hybrid_kv_cache, HybridLinearKVPool): diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 2d3e579bb..e3002d1a5 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -45,7 +45,6 @@ from sglang.srt.mem_cache.radix_cache import ( ) from sglang.srt.mem_cache.utils import convert_to_bigram_key from sglang.srt.observability.metrics_collector import StorageMetricsCollector -from sglang.srt.utils import bind_to_closest_numa_node_cuda if TYPE_CHECKING: from sglang.srt.mem_cache.cache_init_params import CacheInitParams @@ -59,9 +58,6 @@ class HiRadixCache(RadixCache): def __init__(self, params: CacheInitParams, server_args: ServerArgs): self._enable_metrics_flag = params.enable_metrics - if not server_args.disable_hicache_numa_detect: - bind_to_closest_numa_node_cuda() - self.page_size = params.page_size self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d9210c656..b9bd707df 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -555,7 +555,6 @@ class ServerArgs: hicache_write_policy: str = "write_through" hicache_io_backend: str = "kernel" hicache_mem_layout: str = "layer_first" - disable_hicache_numa_detect: bool = False hicache_storage_backend: Optional[str] = None hicache_storage_prefetch_policy: str = "best_effort" hicache_storage_backend_extra_config: Optional[str] = None @@ -5069,11 +5068,6 @@ class ServerArgs: default=ServerArgs.hicache_mem_layout, help="The layout of host memory pool for hierarchical cache.", ) - parser.add_argument( - "--disable-hicache-numa-detect", - action="store_true", - help="Disable binding the process to the NUMA node closest to the active CUDA device when hierarchical cache is enabled.", - ) parser.add_argument( "--hicache-storage-backend", type=str, @@ -5542,7 +5536,7 @@ class ServerArgs: "--numa-node", type=int, nargs="+", - help="Sets the numa node for the subprocesses. i-th element corresponds to i-th subprocess.", + help="Sets the numa node for the subprocesses. i-th element corresponds to i-th subprocess. If unset, will be automatically detected on NUMA systems.", ) parser.add_argument( "--enable-deterministic-inference", diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index a325822ed..e3b737e7e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3408,30 +3408,6 @@ def get_device_sm_nvidia_smi(): return (0, 0) # Default/fallback value -def get_libnuma(): - libnuma = None - - for libnuma_so in ["libnuma.so", "libnuma.so.1"]: - try: - libnuma = ctypes.CDLL(libnuma_so) - except OSError as e: - logger.error(f"{e}") - libnuma = None - if libnuma is not None: - break - return libnuma - - -def numa_bind_to_node(node: int): - libnuma = get_libnuma() - - if libnuma is None or libnuma.numa_available() < 0: - logger.error("numa not available on this system, skip bind action") - else: - libnuma.numa_run_on_node(ctypes.c_int(node)) - libnuma.numa_set_preferred(ctypes.c_int(node)) - - def json_list_type(value): try: return orjson.loads(value) @@ -3729,140 +3705,3 @@ def get_or_create_event_loop(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) return loop - - -def get_numa_node_count() -> int: - """ - Get the number of NUMA nodes available on the system. - Must be called after is_numa_available() is True. - Returns: - int: The number of NUMA nodes. - """ - libnuma = get_libnuma() - return libnuma.numa_max_node() + 1 - - -def is_numa_available() -> bool: - try: - libnuma = get_libnuma() - return libnuma.numa_available() >= 0 - except Exception: - return False - - -def get_system_nvgpu_count() -> int: - """ - Get the total number of GPUs in the system (not affected by CUDA_VISIBLE_DEVICES). - - Returns: - int: The total number of physical GPUs. - """ - result = subprocess.run( - ["nvidia-smi", "--list-gpus"], - capture_output=True, - text=True, - check=True, - ) - gpu_lines = [ - line - for line in result.stdout.strip().split("\n") - if line.strip().startswith("GPU") - ] - return len(gpu_lines) - - -@lru_cache(maxsize=1) -def get_device_numa_node_cuda(gpu_id: int = 0) -> int: - """ - Retrieve the NUMA node ID of the CPU socket closest to the gpu_id. - - First tries to query nvidia-smi topology. If it returns a single NUMA ID, uses that directly. - If it returns multiple NUMA IDs (comma/dash separated), falls back to distributing GPUs - evenly across NUMA nodes based on GPU ID intervals. - - For example, with 8 GPUs and 2 NUMA nodes: GPUs 0-3 -> node 0, GPUs 4-7 -> node 1. - - Returns: - int: The NUMA node ID (e.g., 0, 1). - - Raises: - RuntimeError: If device information cannot be retrieved. - """ - - physical_device_id = get_physical_device_id(gpu_id) - - # Query NUMA topology from nvidia-smi - result = subprocess.run( - ["nvidia-smi", "topo", "-C", "-i", str(physical_device_id)], - capture_output=True, - text=True, - check=True, - ) - - output_line = result.stdout.strip() - prefix = "NUMA IDs of closest CPU:" - - if output_line.startswith(prefix): - numa_id_str = output_line[len(prefix) :].strip() - if numa_id_str.isdigit(): - return int(numa_id_str) - - # Fall back: distribute GPUs evenly across NUMA nodes - numa_count = get_numa_node_count() - gpu_count = get_system_nvgpu_count() - - if gpu_count >= numa_count: - gpus_per_numa = gpu_count // numa_count # >= 1 - numa_node = physical_device_id // gpus_per_numa # 0 ~ numa_count - 1 - else: - logger.warning( - f"GPU count {gpu_count} is less than NUMA count {numa_count}. Using first NUMA node." - ) - numa_node = 0 - - return numa_node - - -def get_numa_node(gpu_id): - numa_node = None - try: - device = get_device() - if device == "cuda": - numa_node = get_device_numa_node_cuda(gpu_id) - else: - logger.info(f"Now only supports NVIDIA devices") - except Exception as e: - logger.error(f"Error: {e}") - - return numa_node - - -@lru_cache(maxsize=1) -def get_current_device_numa_node_cuda() -> int: - """ - Retrieve the NUMA node ID of the CPU socket closest to the currently active CUDA device. - """ - - logical_device_id = torch.cuda.current_device() - numa_node = get_device_numa_node_cuda(logical_device_id) - - return numa_node - - -def nvgpu_available() -> bool: - if not torch.cuda.is_available(): - return False - if torch.version.cuda is None: - return False - return True - - -def bind_to_closest_numa_node_cuda(): - """ - Bind the current process to the NUMA node closest to the active CUDA device. - - Uses `numa` library calls via ctypes to set the CPU affinity of the process. - """ - if is_numa_available() and nvgpu_available(): - node_id = get_current_device_numa_node_cuda() - numa_bind_to_node(node_id) diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index 2c934af0d..5c801679b 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/python/sglang/srt/utils/numa_utils.py @@ -1,26 +1,30 @@ +import ctypes +import glob import logging +import math import multiprocessing import os import random +import shutil import time from contextlib import contextmanager from pathlib import Path +from typing import Optional + +import psutil from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import get_numa_node +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() logger = logging.getLogger(__name__) @contextmanager def configure_subprocess(server_args: ServerArgs, gpu_id: int): - numa_node = None - if (numa_nodes := server_args.numa_node) is not None: - numa_node = numa_nodes[gpu_id] - elif envs.SGLANG_AUTO_NUMA_BIND.get(): - numa_node = get_numa_node(gpu_id) - + numa_node = get_numa_node_if_available(server_args, gpu_id) if numa_node is not None and envs.SGLANG_NUMA_BIND_V2.get(): numactl_args = f"--cpunodebind={numa_node} --membind={numa_node}" executable, debug_str = _create_numactl_executable(numactl_args=numactl_args) @@ -58,3 +62,157 @@ def _mp_set_executable(executable: str, debug_str: str): ), f"{multiprocessing.spawn.get_executable()=}" multiprocessing.spawn.set_executable(old_executable) logger.info(f"mp.set_executable revert to {old_executable}") + + +def get_numa_node_if_available(server_args: ServerArgs, gpu_id: int) -> Optional[int]: + """ + Returns the NUMA node for the given GPU id. If it is not set in the server_args, it will try to query the NUMA node for the GPU. + If the NUMA node is not available, has already been configured externally, or the user lacks permission to set NUMA affinity, it will return None. + + Args: + server_args: The server arguments. + gpu_id: The GPU id. + + Returns: + The NUMA node for the given GPU id or None if it is not available. + """ + if server_args.numa_node is not None: + return server_args.numa_node[gpu_id] + if _is_numa_available(): + queried_numa_node = _query_numa_node_for_gpu(gpu_id) + if len(queried_numa_node) == 0: + return None + if len(queried_numa_node) > 1: + # get_numa_node_for_gpu could return multiple nodes, we use the first one for now. + # I don't think there any hardware configs that would have more than one. + logger.warning( + f"Multiple NUMA nodes found for GPU {gpu_id}: {queried_numa_node}. Using the first one." + ) + return queried_numa_node[0] + return None + + +def get_libnuma(): + libnuma = None + + for libnuma_so in ["libnuma.so", "libnuma.so.1"]: + try: + libnuma = ctypes.CDLL(libnuma_so) + except OSError as e: + logger.debug(f"{e}") + libnuma = None + if libnuma is not None: + break + return libnuma + + +def numa_bind_to_node(node: int): + libnuma = get_libnuma() + + if libnuma is None or libnuma.numa_available() < 0: + logger.warning("numa not available on this system, skip bind action") + else: + libnuma.numa_run_on_node(ctypes.c_int(node)) + libnuma.numa_set_preferred(ctypes.c_int(node)) + + +def _can_set_mempolicy() -> bool: + """Check if the process has permission to use NUMA memory policy syscalls.""" + try: + libnuma = get_libnuma() + if libnuma is None or libnuma.numa_available() < 0: + return False + mode = ctypes.c_int() + ret = libnuma.get_mempolicy( + ctypes.byref(mode), None, ctypes.c_ulong(0), None, ctypes.c_ulong(0) + ) + return ret == 0 + except Exception: + return False + + +def _is_numa_available() -> bool: + """ + Check if NUMA is available and not already configured externally. + """ + if not _is_cuda: + return False + + # Check if this is a numa system. + 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.warning( + "numactl command not found, skipping NUMA node configuration for GPU. Install numactl (e.g., apt-get install numactl) to enable automatic NUMA binding." + ) + return False + + if not _can_set_mempolicy(): + logger.warning( + "User lacks permission to set NUMA affinity, skipping NUMA node configuration for GPU. If using docker, try adding --cap-add SYS_NICE to your docker run command." + ) + return False + + return True + + +def _query_numa_node_for_gpu(device_id: int): + """ + Get the NUMA node affinity list for a GPU device. + + Args: + device_id: GPU device index. + Returns: + List of NUMA node IDs that have affinity with the device. + """ + try: + import pynvml + except ModuleNotFoundError: + logger.warning("pynvml not installed, skipping NUMA node configuration for GPU") + return [] + + try: + pynvml.nvmlInit() + + handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) + numa_node_count = len(glob.glob("/sys/devices/system/node/node[0-9]*")) + + c_ulong_bits = ctypes.sizeof(ctypes.c_ulong) * 8 + node_set_size = max(1, math.ceil(numa_node_count / c_ulong_bits)) + node_set = pynvml.nvmlDeviceGetMemoryAffinity( + handle, + node_set_size, + pynvml.NVML_AFFINITY_SCOPE_NODE, + ) + + # Decode the bitmask into a list of NUMA node IDs + numa_nodes = [] + for node_id in range(numa_node_count): + mask_array_index = node_id // c_ulong_bits + mask_bit_index = node_id % c_ulong_bits + if node_set[mask_array_index] & (1 << mask_bit_index): + numa_nodes.append(node_id) + return numa_nodes + except pynvml.NVMLError as e: + logger.warning( + f"NVML error querying memory affinity for GPU {device_id}: {e}, skipping NUMA node configuration for GPU" + ) + return [] + finally: + try: + pynvml.nvmlShutdown() + except Exception: + pass # Ignore shutdown errors diff --git a/test/registered/utils/test_numa_utils.py b/test/registered/utils/test_numa_utils.py new file mode 100644 index 000000000..01209292e --- /dev/null +++ b/test/registered/utils/test_numa_utils.py @@ -0,0 +1,311 @@ +import ctypes +import unittest +from unittest.mock import MagicMock, patch + +from sglang.srt.utils.numa_utils import ( + _is_numa_available, + _query_numa_node_for_gpu, + get_numa_node_if_available, +) +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci + +register_cpu_ci(est_time=1, suite="stage-a-cpu-only") +register_cuda_ci(est_time=10, suite="stage-c-test-4-gpu-gb200") +register_cuda_ci(est_time=10, suite="stage-c-test-8-gpu-b200") + + +class TestIsNumaAvailable(unittest.TestCase): + """Tests for _is_numa_available on both NUMA and non-NUMA systems.""" + + @patch("sglang.srt.utils.numa_utils._is_cuda", False) + def test_returns_false_when_not_cuda(self): + self.assertFalse(_is_numa_available()) + + @patch("sglang.srt.utils.numa_utils._is_cuda", True) + @patch("os.path.isdir", return_value=False) + 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 + ): + 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 + ): + 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 + ): + 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") + + +class TestQueryNumaNodeForGpu(unittest.TestCase): + """Tests for _query_numa_node_for_gpu with mocked pynvml.""" + + @patch( + "sglang.srt.utils.numa_utils.glob.glob", + return_value=[ + "/sys/devices/system/node/node0", + "/sys/devices/system/node/node1", + ], + ) + def test_single_node_affinity(self, _mock_glob): + c_ulong_bits = ctypes.sizeof(ctypes.c_ulong) * 8 + # Bitmask: bit 0 set -> node 0 + node_set = [1] + + mock_pynvml = MagicMock() + mock_pynvml.nvmlDeviceGetMemoryAffinity.return_value = node_set + mock_pynvml.NVML_AFFINITY_SCOPE_NODE = 0 + + with patch.dict("sys.modules", {"pynvml": mock_pynvml}): + result = _query_numa_node_for_gpu(0) + + self.assertEqual(result, [0]) + mock_pynvml.nvmlInit.assert_called_once() + mock_pynvml.nvmlShutdown.assert_called_once() + + @patch( + "sglang.srt.utils.numa_utils.glob.glob", + return_value=[ + "/sys/devices/system/node/node0", + "/sys/devices/system/node/node1", + ], + ) + def test_second_node_affinity(self, _mock_glob): + # Bitmask: bit 1 set -> node 1 + node_set = [2] + + mock_pynvml = MagicMock() + mock_pynvml.nvmlDeviceGetMemoryAffinity.return_value = node_set + mock_pynvml.NVML_AFFINITY_SCOPE_NODE = 0 + + with patch.dict("sys.modules", {"pynvml": mock_pynvml}): + result = _query_numa_node_for_gpu(1) + + self.assertEqual(result, [1]) + + @patch( + "sglang.srt.utils.numa_utils.glob.glob", + return_value=[ + "/sys/devices/system/node/node0", + "/sys/devices/system/node/node1", + "/sys/devices/system/node/node2", + "/sys/devices/system/node/node3", + ], + ) + def test_multiple_node_affinity(self, _mock_glob): + # Bitmask: bits 1 and 3 set -> nodes 1, 3 (binary: ...1010 = 10) + node_set = [0b1010] + + mock_pynvml = MagicMock() + mock_pynvml.nvmlDeviceGetMemoryAffinity.return_value = node_set + mock_pynvml.NVML_AFFINITY_SCOPE_NODE = 0 + + with patch.dict("sys.modules", {"pynvml": mock_pynvml}): + result = _query_numa_node_for_gpu(0) + + self.assertEqual(result, [1, 3]) + + @patch( + "sglang.srt.utils.numa_utils.glob.glob", + return_value=[ + "/sys/devices/system/node/node0", + "/sys/devices/system/node/node1", + ], + ) + def test_no_affinity(self, _mock_glob): + node_set = [0] + + mock_pynvml = MagicMock() + mock_pynvml.nvmlDeviceGetMemoryAffinity.return_value = node_set + mock_pynvml.NVML_AFFINITY_SCOPE_NODE = 0 + + with patch.dict("sys.modules", {"pynvml": mock_pynvml}): + result = _query_numa_node_for_gpu(0) + + self.assertEqual(result, []) + + @patch( + "sglang.srt.utils.numa_utils.glob.glob", + return_value=[ + "/sys/devices/system/node/node0", + "/sys/devices/system/node/node1", + ], + ) + def test_nvml_shutdown_called_on_success(self, _mock_glob): + node_set = [1] + mock_pynvml = MagicMock() + mock_pynvml.nvmlDeviceGetMemoryAffinity.return_value = node_set + mock_pynvml.NVML_AFFINITY_SCOPE_NODE = 0 + + with patch.dict("sys.modules", {"pynvml": mock_pynvml}): + _query_numa_node_for_gpu(0) + + mock_pynvml.nvmlShutdown.assert_called_once() + + +class TestGetNumaNodeIfAvailable(unittest.TestCase): + """Tests for get_numa_node_if_available combining _is_numa_available + _query_numa_node_for_gpu.""" + + def _make_server_args(self, numa_node=None): + args = MagicMock() + args.numa_node = numa_node + return args + + def test_returns_explicit_numa_node_from_server_args(self): + args = self._make_server_args(numa_node=[2, 3, 0, 1]) + self.assertEqual(get_numa_node_if_available(args, 0), 2) + self.assertEqual(get_numa_node_if_available(args, 1), 3) + self.assertEqual(get_numa_node_if_available(args, 2), 0) + self.assertEqual(get_numa_node_if_available(args, 3), 1) + + @patch("sglang.srt.utils.numa_utils._is_numa_available", return_value=False) + def test_returns_none_when_numa_not_available(self, _mock_avail): + args = self._make_server_args(numa_node=None) + self.assertIsNone(get_numa_node_if_available(args, 0)) + + @patch("sglang.srt.utils.numa_utils._query_numa_node_for_gpu", return_value=[]) + @patch("sglang.srt.utils.numa_utils._is_numa_available", return_value=True) + def test_returns_none_when_query_returns_empty(self, _mock_avail, _mock_gpu): + args = self._make_server_args(numa_node=None) + self.assertIsNone(get_numa_node_if_available(args, 0)) + + @patch("sglang.srt.utils.numa_utils._query_numa_node_for_gpu", return_value=[1]) + @patch("sglang.srt.utils.numa_utils._is_numa_available", return_value=True) + def test_returns_queried_single_node(self, _mock_avail, _mock_gpu): + args = self._make_server_args(numa_node=None) + self.assertEqual(get_numa_node_if_available(args, 0), 1) + + @patch("sglang.srt.utils.numa_utils._query_numa_node_for_gpu", return_value=[0, 2]) + @patch("sglang.srt.utils.numa_utils._is_numa_available", return_value=True) + def test_returns_first_node_when_multiple_found(self, _mock_avail, _mock_gpu): + args = self._make_server_args(numa_node=None) + self.assertEqual(get_numa_node_if_available(args, 0), 0) + + @patch("sglang.srt.utils.numa_utils._query_numa_node_for_gpu", return_value=[0, 2]) + @patch("sglang.srt.utils.numa_utils._is_numa_available", return_value=True) + def test_logs_warning_when_multiple_nodes(self, _mock_avail, _mock_gpu): + args = self._make_server_args(numa_node=None) + with self.assertLogs("sglang.srt.utils.numa_utils", level="WARNING") as cm: + get_numa_node_if_available(args, 0) + self.assertTrue(any("Multiple NUMA nodes" in msg for msg in cm.output)) + + @patch("sglang.srt.utils.numa_utils._is_numa_available", return_value=True) + @patch("sglang.srt.utils.numa_utils._query_numa_node_for_gpu", return_value=[1]) + def test_explicit_server_args_takes_precedence(self, _mock_gpu, _mock_avail): + args = self._make_server_args(numa_node=[5, 6]) + result = get_numa_node_if_available(args, 0) + self.assertEqual(result, 5) + _mock_avail.assert_not_called() + _mock_gpu.assert_not_called() + + +def _get_gpu_name(): + try: + import pynvml + + pynvml.nvmlInit() + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + name = pynvml.nvmlDeviceGetName(handle) + pynvml.nvmlShutdown() + return name + except Exception: + return "" + + +_gpu_name = _get_gpu_name() + + +@unittest.skipUnless("GB200" in _gpu_name, "Requires GB200 hardware") +class TestGB200NumaTopology(unittest.TestCase): + """Hardware test validating expected NUMA topology on GB200 (2 NUMA nodes, 4 GPUs).""" + + def _make_server_args(self): + args = MagicMock() + args.numa_node = None + return args + + def test_gpu_numa_mapping(self): + expected = {0: 0, 1: 0, 2: 1, 3: 1} + args = self._make_server_args() + for gpu_id, expected_node in expected.items(): + result = get_numa_node_if_available(args, gpu_id) + self.assertEqual( + result, + expected_node, + f"GPU {gpu_id}: expected NUMA node {expected_node}, got {result}", + ) + + +@unittest.skipUnless("B200" in _gpu_name, "Requires B200 hardware") +class TestB200NumaTopology(unittest.TestCase): + """Hardware test validating expected NUMA topology on B200 (2 NUMA nodes, 8 GPUs).""" + + def _make_server_args(self): + args = MagicMock() + args.numa_node = None + return args + + def test_gpu_numa_mapping(self): + expected = {0: 0, 1: 0, 2: 0, 3: 0, 4: 1, 5: 1, 6: 1, 7: 1} + args = self._make_server_args() + for gpu_id, expected_node in expected.items(): + result = get_numa_node_if_available(args, gpu_id) + self.assertEqual( + result, + expected_node, + f"GPU {gpu_id}: expected NUMA node {expected_node}, got {result}", + ) + + +if __name__ == "__main__": + unittest.main()