chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

35
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -46,6 +46,7 @@ from sglang.srt.utils import (
configure_logger,
get_zmq_socket,
kill_itself_when_parent_died,
maybe_reindex_device_id,
)
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
@@ -139,6 +140,9 @@ class DataParallelController:
# Load balance budget
self.dp_budget = DPBudget()
# To protect changing env vars to set CUDA_VISIBLE_DEVICES.
self.env_lock = threading.Lock()
# Launch data parallel workers
self.scheduler_procs = []
self.workers: List[zmq.Socket] = [None] * server_args.dp_size
@@ -399,21 +403,22 @@ class DataParallelController:
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
proc = mp.Process(
target=run_scheduler_process,
args=(
server_args,
rank_port_args,
gpu_id,
tp_rank,
moe_ep_rank,
pp_rank,
dp_rank,
writer,
),
)
with memory_saver_adapter.configure_subprocess():
proc.start()
with self.env_lock, maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
target=run_scheduler_process,
args=(
server_args,
rank_port_args,
gpu_id,
tp_rank,
moe_ep_rank,
pp_rank,
dp_rank,
writer,
),
)
with memory_saver_adapter.configure_subprocess():
proc.start()
self.scheduler_procs.append(proc)
scheduler_pipe_readers.append(reader)