From a8f12a63fde4765dffe53f7bf1482d52ac80af33 Mon Sep 17 00:00:00 2001 From: Richard Liu <39319471+richardsliu@users.noreply.github.com> Date: Thu, 20 Mar 2025 07:59:33 -0700 Subject: [PATCH] Fix env vars for running Ray distributed backend on GKE (#15166) Signed-off-by: Richard Liu --- vllm/executor/ray_distributed_executor.py | 2 ++ vllm/platforms/interface.py | 2 ++ vllm/platforms/tpu.py | 4 ++++ 3 files changed, 8 insertions(+) diff --git a/vllm/executor/ray_distributed_executor.py b/vllm/executor/ray_distributed_executor.py index 18ff32155c5fc..d769d235020d5 100644 --- a/vllm/executor/ray_distributed_executor.py +++ b/vllm/executor/ray_distributed_executor.py @@ -340,6 +340,8 @@ class RayDistributedExecutor(DistributedExecutorBase): and v not in self.non_carry_over_env_vars ] + env_vars_to_copy.extend(current_platform.additional_env_vars) + # Copy existing env vars to each worker's args for args in all_args_to_update_environment_variables: # TODO: refactor platform-specific env vars diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7415b5d5f0603..c7152d0bfb79e 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -112,6 +112,8 @@ class Platform: supported_quantization: list[str] = [] + additional_env_vars: list[str] = [] + def is_cuda(self) -> bool: return self._enum == PlatformEnum.CUDA diff --git a/vllm/platforms/tpu.py b/vllm/platforms/tpu.py index 8e2c28d9327b5..073d46c25d579 100644 --- a/vllm/platforms/tpu.py +++ b/vllm/platforms/tpu.py @@ -29,6 +29,10 @@ class TpuPlatform(Platform): "tpu_int8", "compressed-tensors", "compressed_tensors" ] + additional_env_vars: list[str] = [ + "TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS" + ] + @classmethod def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int, dtype: torch.dtype, kv_cache_dtype: Optional[str],