[sglang-miles] Cherry-pick #24766: Fix NUMA NVML handle resolution under shuffled CUDA_VISIBLE_DEVICES

python/sglang/srt/utils/numa_utils.py

@@ -12,6 +12,7 @@
 from typing import Optional
 
 import psutil
+import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.server_args import ServerArgs
@@ -24,14 +25,22 @@
 
 @contextmanager
 def configure_subprocess(server_args: ServerArgs, gpu_id: int):
-    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)
-        with _mp_set_executable(executable=executable, debug_str=debug_str):
-            yield
-    else:
-        yield
+    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
+    yield
 
 
 def _create_numactl_executable(numactl_args: str):
@@ -64,6 +73,21 @@ def _mp_set_executable(executable: str, debug_str: str):
         logger.info(f"mp.set_executable revert to {old_executable}")
 
 
+def _get_nvml_device_index(device_id: int) -> int:
+    # _get_nvml_device_index is an internal PyTorch helper, so fall back to
+    # device_id directly if the helper is unavailable.
+    get_nvml_device_index = getattr(torch.cuda, "_get_nvml_device_index", None)
+    if get_nvml_device_index is None:
+        logger.warning(
+            "torch.cuda._get_nvml_device_index is unavailable; falling back to "
+            f"device_id={device_id} as the NVML device index. This may select "
+            "the wrong physical GPU when CUDA_VISIBLE_DEVICES reorders devices "
+            f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')})."
+        )
+        return device_id
+    return get_nvml_device_index(device_id)
+
+
 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.
@@ -174,7 +198,7 @@ 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.
+        device_id: CUDA logical device index (post-CUDA_VISIBLE_DEVICES).
     Returns:
         List of NUMA node IDs that have affinity with the device.
     """
@@ -187,7 +211,11 @@ def _query_numa_node_for_gpu(device_id: int):
     try:
         pynvml.nvmlInit()
 
-        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
+        # device_id is a CUDA logical index. Convert it to the corresponding
+        # NVML index so reordered CUDA_VISIBLE_DEVICES maps to the right GPU.
+        # _get_nvml_device_index takes CUDA_VISIBLE_DEVICES into account.
+        nvml_device_id = _get_nvml_device_index(device_id)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(nvml_device_id)
         numa_node_count = len(glob.glob("/sys/devices/system/node/node[0-9]*"))
 
         c_ulong_bits = ctypes.sizeof(ctypes.c_ulong) * 8