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, 21:00 UTC
[bugfix] Fix profiling for RayDistributedExecutor (#13945)
Signed-off-by: Rui Qiao <ruisearch42@gmail.com>
Rui Qiao committed
Feb 27, 2025 at 09:01 UTC
06c8f8d885738a6525fd6774186f094e15676054
1 file changed
+13
-8
vllm/executor/ray_distributed_executor.py
+13
-8
@@ -309,19 +309,24 @@ class RayDistributedExecutor(DistributedExecutorBase):
309
",".join(map(str, node_gpus[node_id])),
310
} for (node_id, _) in worker_node_and_gpu_ids]
311
312
+ # Environment variables to copy from driver to workers
313
+ env_vars_to_copy = [
314
+ "VLLM_ATTENTION_BACKEND", "TPU_CHIPS_PER_HOST_BOUNDS",
315
+ "TPU_HOST_BOUNDS", "VLLM_USE_V1", "VLLM_TRACE_FUNCTION",
316
+ "VLLM_TORCH_PROFILER_DIR", "VLLM_TEST_ENABLE_EP"
317
+ ]
318
+
319
+ # Copy existing env vars to each worker's args
320
for args in all_args_to_update_environment_variables:
313
- # some carry-over env vars from the driver
321
# TODO: refactor platform-specific env vars
315
- for name in [
316
- "VLLM_ATTENTION_BACKEND",
317
- "TPU_CHIPS_PER_HOST_BOUNDS",
318
- "TPU_HOST_BOUNDS",
319
- "VLLM_USE_V1",
320
- "VLLM_TRACE_FUNCTION",
321
- ]:
322
+ for name in env_vars_to_copy:
323
if name in os.environ:
324
args[name] = os.environ[name]
325
326
+ logger.info(
327
+ "Copying the following environment variables to workers: %s",
328
+ [v for v in env_vars_to_copy if v in os.environ])
329
+
330
self._env_vars_for_all_workers = (
331
all_args_to_update_environment_variables)
332