diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index d0a417a788..cd94138654 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -15,6 +15,7 @@ import pytest from tests.utils import multi_gpu_test from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.executor import multiproc_executor from vllm.v1.executor.multiproc_executor import MultiprocExecutor @@ -378,8 +379,8 @@ def test_multiproc_executor_properties(): def test_multiproc_executor_multi_node(): """Test MultiprocExecutor with multi-node configuration. This simulates 2 nodes with TP=4: - - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2 - - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2 + - Node 0 (rank 0): Uses GPUs 0,1 (device_control_env_var=0,1) with TP=2 + - Node 1 (rank 1): Uses GPUs 2,3 (device_control_env_var=2,3) with TP=2 Total world_size = 4, nnodes = 2 """ # Python 3.14+ changed default multiprocessing start method to 'forkserver' @@ -392,15 +393,20 @@ def test_multiproc_executor_multi_node(): # symm_mem does not work for simulating multi instance in single node os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" + # Use the platform-specific device control env var (e.g. + # CUDA_VISIBLE_DEVICES on CUDA, ZE_AFFINITY_MASK on XPU) so that each + # node only sees its own devices on different platform. + device_control_env_var = current_platform.device_control_env_var + def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int): """Run a single node's executor.""" executor = None try: - # Set CUDA_VISIBLE_DEVICES for this node + # Restrict visible devices for this node if node_rank == 0: - os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" + os.environ[device_control_env_var] = "0,1" else: - os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" + os.environ[device_control_env_var] = "2,3" # Create config for this node vllm_config = create_vllm_config(