add reindex_device_id to device OOT plugin (#36220)

Co-authored-by: Akash Palla <apalla@fb.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/sglang/srt/platforms/cuda.py

@@ -1,5 +1,8 @@
 """CUDA device operations for the SRT platform layer."""
 
+import logging
+import os
+from contextlib import contextmanager
 from typing import Optional
 
 import torch
@@ -11,6 +14,8 @@ from sglang.srt.platforms.device_mixin import (
 )
 from sglang.srt.platforms.interface import SRTPlatform
 
+logger = logging.getLogger(__name__)
+
 
 class CudaDeviceMixin(DeviceMixin):
     """CUDA implementation of the shared device operations."""
@@ -57,6 +62,34 @@ class CudaDeviceMixin(DeviceMixin):
             return False
         return True
 
+    @contextmanager
+    def reindex_device_id(self, gpu_id: int):
+        if os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID":
+            logger.warning(
+                "`CUDA_DEVICE_ORDER` is not set to `PCI_BUS_ID`. Please set "
+                "`CUDA_DEVICE_ORDER=PCI_BUS_ID` to avoid unexpected behavior."
+            )
+
+        original_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+        if original_cuda_visible_devices:
+            cuda_visible_devices = original_cuda_visible_devices.split(",")
+        else:
+            cuda_visible_devices = []
+
+        str_gpu_id = (
+            cuda_visible_devices[gpu_id] if cuda_visible_devices else str(gpu_id)
+        )
+        os.environ["CUDA_VISIBLE_DEVICES"] = str_gpu_id
+
+        logger.debug(f"Set CUDA_VISIBLE_DEVICES to {str_gpu_id}")
+
+        yield 0
+
+        if original_cuda_visible_devices:
+            os.environ["CUDA_VISIBLE_DEVICES"] = original_cuda_visible_devices
+        else:
+            del os.environ["CUDA_VISIBLE_DEVICES"]
+
     def get_torch_distributed_backend_str(self) -> str:
         return "nccl"
 

python/sglang/srt/platforms/device_mixin.py

@@ -27,6 +27,8 @@ Method status annotations:
 
 import enum
 import random
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import NamedTuple, Optional
 
 import numpy as np
@@ -163,6 +165,11 @@ class DeviceMixin:
         """[Active] Whether pinned host memory is available for a target device."""
         return False
 
+    @contextmanager
+    def reindex_device_id(self, device_id: int) -> Iterator[int]:
+        """[Active] Temporarily remap a physical device to logical device 0."""
+        yield device_id
+
     # ------------------------------------------------------------------
     # Planned methods — reserved interface.  Core still uses hardcoded
     # calls (e.g. torch.cuda.*).  OOT implementations will NOT take

python/sglang/srt/utils/common.py

@@ -65,6 +65,7 @@ from typing import (
     Callable,
     Dict,
     Generic,
+    Iterator,
     List,
     NamedTuple,
     Optional,
@@ -1171,29 +1172,13 @@ def get_device_sm_nvidia_smi():
 
 
 @contextmanager
-def maybe_reindex_device_id(gpu_id: int):
-
-    if envs.SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS.get() is False or not is_cuda_alike():
+def maybe_reindex_device_id(gpu_id: int) -> Iterator[int]:
+    if not envs.SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS.get():
         yield gpu_id
         return
 
-    original_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
-    if original_cuda_visible_devices:
-        cuda_visible_devices = original_cuda_visible_devices.split(",")
-    else:
-        cuda_visible_devices = []
-
-    str_gpu_id = cuda_visible_devices[gpu_id] if cuda_visible_devices else str(gpu_id)
-    os.environ["CUDA_VISIBLE_DEVICES"] = str_gpu_id
-
-    logger.debug(f"Set CUDA_VISIBLE_DEVICES to {str_gpu_id}")
-
-    yield 0
-
-    if original_cuda_visible_devices:
-        os.environ["CUDA_VISIBLE_DEVICES"] = original_cuda_visible_devices
-    else:
-        del os.environ["CUDA_VISIBLE_DEVICES"]
+    with current_platform.reindex_device_id(gpu_id) as reindexed_device_id:
+        yield reindexed_device_id
 
 
 cached_device_index = -1