A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
Fix env vars for running Ray distributed backend on GKE (#15166)
Signed-off-by: Richard Liu <ricliu@google.com>
Richard Liu committed
Mar 20, 2025 at 07:59 UTC
a8f12a63fde4765dffe53f7bf1482d52ac80af33
3 files changed
+8
vllm/executor/ray_distributed_executor.py
+2
@@ -340,6 +340,8 @@ class RayDistributedExecutor(DistributedExecutorBase):
340
and v not in self.non_carry_over_env_vars
341
]
342
343
+ env_vars_to_copy.extend(current_platform.additional_env_vars)
344
+
345
# Copy existing env vars to each worker's args
346
for args in all_args_to_update_environment_variables:
347
# TODO: refactor platform-specific env vars
vllm/platforms/interface.py
+2
@@ -112,6 +112,8 @@ class Platform:
112
113
supported_quantization: list[str] = []
114
115
+ additional_env_vars: list[str] = []
116
+
117
def is_cuda(self) -> bool:
118
return self._enum == PlatformEnum.CUDA
119
vllm/platforms/tpu.py
+4
@@ -29,6 +29,10 @@ class TpuPlatform(Platform):
29
"tpu_int8", "compressed-tensors", "compressed_tensors"
30
]
31
32
+ additional_env_vars: list[str] = [
33
+ "TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS"
34
+ ]
35
+
36
@classmethod
37
def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
38
dtype: torch.dtype, kv_cache_dtype: Optional[str],