[Core] replace narrow-usage RayWorkerVllm to general WorkerWrapper to reduce code duplication

tests/distributed/test_pynccl.py

@@ -6,6 +6,7 @@
 
 from vllm.distributed.device_communicators.pynccl import (NCCLCommunicator,
                                                           ncclGetUniqueId)
+from vllm.utils import update_environment_variables
 
 
 def distributed_run(fn, world_size):
@@ -32,8 +33,7 @@ def update_env(fn):
     # so we need to pass the environment variables as arguments
     # and update the environment variables in the function
     def wrapper(env):
-        import os
-        os.environ.update(env)
+        update_environment_variables(env)
         fn()
 
     return wrapper

vllm/engine/ray_utils.py

@@ -1,9 +1,9 @@
 import pickle
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from vllm.config import ParallelConfig
 from vllm.logger import init_logger
-from vllm.utils import get_ip, is_hip, set_cuda_visible_devices
+from vllm.utils import get_ip, is_hip, update_environment_variables
 
 logger = init_logger(__name__)
 
@@ -52,8 +52,8 @@ def get_node_and_gpu_ids(self) -> Tuple[str, List[int]]:
             gpu_ids = ray.get_gpu_ids()
             return node_id, gpu_ids
 
-        def set_cuda_visible_devices(self, device_ids) -> None:
-            set_cuda_visible_devices(device_ids)
+        def update_environment_variables(self, envs: Dict[str, str]) -> None:
+            update_environment_variables(envs)
 
         def execute_model_compiled_dag_remote(self, ignored):
             """Used only when compiled DAG is enabled."""

vllm/executor/ray_gpu_executor.py

@@ -14,7 +14,7 @@
 from vllm.lora.request import LoRARequest
 from vllm.sequence import SamplerOutput, SequenceGroupMetadata
 from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
-                        make_async, set_cuda_visible_devices)
+                        make_async, update_environment_variables)
 
 if ray is not None:
     from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
@@ -134,9 +134,16 @@ def _init_workers_ray(self, placement_group: "PlacementGroup",
             node_gpus[node_id] = sorted(gpu_ids)
 
         # Set CUDA_VISIBLE_DEVICES for the driver and workers.
-        set_cuda_visible_devices(node_gpus[driver_node_id])
+        # ",".join(map(str, device_ids))
+        update_environment_variables({
+            "CUDA_VISIBLE_DEVICES":
+            ",".join(map(str, node_gpus[driver_node_id]))
+        })
         for worker, (node_id, _) in zip(self.workers, worker_node_and_gpu_ids):
-            worker.set_cuda_visible_devices.remote(node_gpus[node_id])
+            worker.update_environment_variables.remote({
+                "CUDA_VISIBLE_DEVICES":
+                ",".join(map(str, node_gpus[node_id]))
+            })
 
         distributed_init_method = get_distributed_init_method(
             driver_ip, get_open_port())

vllm/utils.py

@@ -232,8 +232,14 @@ def get_open_port() -> int:
             return s.getsockname()[1]
 
 
-def set_cuda_visible_devices(device_ids: List[int]) -> None:
-    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, device_ids))
+def update_environment_variables(envs: Dict[str, str]):
+    for k, v in envs.items():
+        if k in os.environ:
+            warnings.warn(
+                f"Overwriting environment variable {k} "
+                f"from {os.environ[k]} to {v}",
+                stacklevel=2)
+        os.environ[k] = v
 
 
 def chunk_list(lst, chunk_size):