Set CUDA_VISIBLE_DEVICES to achieve one GPU per process (#9170)

Co-authored-by: SangBin Cho <rkooo567@gmail.com>
Co-authored-by: Cheng Wan <cwan@x.ai>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2025-10-17 17:30:06 -07:00
committed by GitHub
parent 69fe3c9726
commit 9eefe2c0b7
7 changed files with 69 additions and 51 deletions

View File

@@ -88,6 +88,7 @@ from torch.profiler import ProfilerActivity, profile, record_function
from torch.utils._contextlib import _DecoratorContextManager
from typing_extensions import Literal
from sglang.srt.environ import envs
from sglang.srt.metrics.func_timer import enable_func_timer
logger = logging.getLogger(__name__)
@@ -3273,7 +3274,7 @@ def json_list_type(value):
@contextmanager
def maybe_reindex_device_id(gpu_id: int):
if not is_cuda_alike():
if envs.SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS.get() is False or not is_cuda_alike():
yield gpu_id
return