[XPU][UT]Using the device_control_env_var to restrict visible device on different platform (#57581)
Signed-off-by: yisheng <yi.sheng@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
YiSheng5 committed
Sep 20, 2026 at 10:52 UTC
f4fc5aae7fea790e78bb2b5ed62f636dc20c71d8
1 file changed
+11
-5
tests/distributed/test_multiproc_executor.py
+11
-5
@@ -15,6 +15,7 @@ import pytest
15
from tests.utils import multi_gpu_test
16
from vllm.config import VllmConfig
17
from vllm.engine.arg_utils import EngineArgs
18
+from vllm.platforms import current_platform
19
from vllm.v1.core.sched.output import SchedulerOutput
20
from vllm.v1.executor import multiproc_executor
21
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
@@ -378,8 +379,8 @@ def test_multiproc_executor_properties():
379
def test_multiproc_executor_multi_node():
380
"""Test MultiprocExecutor with multi-node configuration.
381
This simulates 2 nodes with TP=4:
381
- - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2
382
- - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
382
+ - Node 0 (rank 0): Uses GPUs 0,1 (device_control_env_var=0,1) with TP=2
383
+ - Node 1 (rank 1): Uses GPUs 2,3 (device_control_env_var=2,3) with TP=2
384
Total world_size = 4, nnodes = 2
385
"""
386
# Python 3.14+ changed default multiprocessing start method to 'forkserver'
@@ -392,15 +393,20 @@ def test_multiproc_executor_multi_node():
393
# symm_mem does not work for simulating multi instance in single node
394
os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
395
396
+ # Use the platform-specific device control env var (e.g.
397
+ # CUDA_VISIBLE_DEVICES on CUDA, ZE_AFFINITY_MASK on XPU) so that each
398
+ # node only sees its own devices on different platform.
399
+ device_control_env_var = current_platform.device_control_env_var
400
+
401
def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int):
402
"""Run a single node's executor."""
403
executor = None
404
try:
399
- # Set CUDA_VISIBLE_DEVICES for this node
405
+ # Restrict visible devices for this node
406
if node_rank == 0:
401
- os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
407
+ os.environ[device_control_env_var] = "0,1"
408
else:
403
- os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
409
+ os.environ[device_control_env_var] = "2,3"
410
411
# Create config for this node
412
vllm_config = create_vllm_config(