diff --git a/python/sglang/srt/platforms/cuda.py b/python/sglang/srt/platforms/cuda.py index d1e9f5e24..3ae863561 100644 --- a/python/sglang/srt/platforms/cuda.py +++ b/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" diff --git a/python/sglang/srt/platforms/device_mixin.py b/python/sglang/srt/platforms/device_mixin.py index fd51a24ae..f3369f185 100644 --- a/python/sglang/srt/platforms/device_mixin.py +++ b/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 diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 04b04eb9c..35614abcb 100644 --- a/python/sglang/srt/utils/common.py +++ b/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