Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 28 additions & 15 deletions tests/test_vllm_generate_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from dataclasses import dataclass

import vime.utils.external_utils.command_utils as U
from vime.backends.vllm_utils.vllm_engine import _wait_server_healthy, launch_server_process
from vime.backends.vllm_utils.vllm_engine import _compute_server_args, _wait_server_healthy, launch_server_process
from vime.rollout import vllm_rollout
from vime.utils import http_utils
from vime.utils.types import Sample
Expand Down Expand Up @@ -65,13 +65,6 @@ def _free_port() -> int:
return sock.getsockname()[1]


def _visible_devices(num_gpus: int) -> str:
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
if visible_devices:
return ",".join(visible_devices.split(",")[:num_gpus])
return ",".join(str(i) for i in range(num_gpus))


def _stop_process_tree(process) -> None:
if not process.is_alive():
return
Expand Down Expand Up @@ -113,24 +106,44 @@ def _execute_case(case: VLLMGenerateCase):
U.exec_command(f"hf download {case.hf_repo} --local-dir {case.model_path}")

server_port = _free_port()
server_args = Namespace(
args = Namespace(
rollout_num_gpus_per_engine=case.num_gpus,
num_gpus_per_node=case.num_gpus,
hf_checkpoint=case.model_path,
seed=1234,
vllm_gpu_memory_utilization=0.9,
vllm_async_scheduling=False,
vllm_enforce_eager=False,
vllm_enable_sleep_mode=False,
rollout_max_context_len=case.max_model_len,
use_rollout_routing_replay=case.use_rollout_routing_replay,
vllm_dp_size=1,
# Placement attrs so _compute_server_args derives the GPU base through the
# real get_base_gpu_id() path rather than a hardcoded id. A single
# colocate engine at rank 0 -> local base 0, and the child server's
# CUDA_VISIBLE_DEVICES is computed exactly as the production VLLMEngine
# does (honoring an externally set CUDA_VISIBLE_DEVICES instead of always
# grabbing physical GPU 0).
colocate=True,
actor_num_nodes=1,
actor_num_gpus_per_node=case.num_gpus,
use_critic=False,
debug_rollout_only=False,
)

process = launch_server_process(
bind_host="127.0.0.1",
server_port=server_port,
args=server_args,
# ``launch_server_process`` consumes a single ``server_args`` dict built by
# ``_compute_server_args`` since the PR #68 multi-node topology refactor
# (this test predates it). Let it derive GPU placement via
# get_base_gpu_id()/_to_local_gpu_id() — identical to VLLMEngine — so the
# launched server tracks CUDA_VISIBLE_DEVICES.
server_args = _compute_server_args(
args,
rank=0,
visible_devices=_visible_devices(case.num_gpus),
model_path=case.model_path,
dist_init_addr=None,
host="127.0.0.1",
port=server_port,
)
process = launch_server_process(server_args)
Comment on lines +139 to +146

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

By removing the _visible_devices helper and relying solely on _compute_server_args, the test inherits a limitation in _compute_server_args (specifically in _to_local_gpu_id). When CUDA_VISIBLE_DEVICES is set to non-zero/non-consecutive indices (e.g., "2,3" in a multi-GPU or containerized environment), _to_local_gpu_id maps the GPU ID to a local index (e.g., 0), resulting in server_args["visible_devices"] being set to "0". When the vllm serve subprocess is launched, its CUDA_VISIBLE_DEVICES is set to "0", causing it to attempt to run on physical GPU 0 instead of the allocated physical GPU 2. This can lead to CUDA initialization failures or resource conflicts in multi-GPU environments. To prevent this, we should explicitly override server_args["visible_devices"] with the correct physical GPU IDs from CUDA_VISIBLE_DEVICES.

    server_args = _compute_server_args(
        args,
        rank=0,
        dist_init_addr=None,
        host="127.0.0.1",
        port=server_port,
        base_gpu_id=0,
    )
    cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
    if cvd:
        server_args["visible_devices"] = ",".join(cvd.split(",")[:case.num_gpus])
    else:
        server_args["visible_devices"] = ",".join(str(i) for i in range(case.num_gpus))
    process = launch_server_process(server_args)


try:
_wait_server_healthy(f"http://127.0.0.1:{server_port}", process)
Expand Down
Loading