diff --git a/tests/distributed/test_multi_node_assignment.py b/tests/distributed/test_multi_node_assignment.py index 5d3f524f4d..6346272429 100644 --- a/tests/distributed/test_multi_node_assignment.py +++ b/tests/distributed/test_multi_node_assignment.py @@ -11,6 +11,7 @@ pytest distributed/test_multi_node_assignment.py """ import os +import time import pytest import ray @@ -19,7 +20,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from vllm import initialize_ray_cluster from vllm.config import ParallelConfig from vllm.utils.network_utils import get_ip -from vllm.v1.executor.ray_utils import _wait_until_pg_removed +from vllm.v1.executor.ray_utils import available_resources_per_node VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1" @@ -61,4 +62,18 @@ def test_multi_node_assignment() -> None: for worker in workers: ray.kill(worker) - _wait_until_pg_removed(config.placement_group) + ray.util.remove_placement_group(config.placement_group) + + gpus_needed = len(workers) + deadline = time.monotonic() + 60 + node_id = ray.get_runtime_context().get_node_id() + while True: + free_gpus = available_resources_per_node().get(node_id, {}).get("GPU", 0) + if free_gpus >= gpus_needed: + break + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("placement group resources were not released") + + time.sleep(min(1, remaining)) diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 071a8eb88c..305d97368a 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -517,24 +517,6 @@ def _wait_until_pg_ready(current_placement_group: "PlacementGroup"): ) from None -def _wait_until_pg_removed(current_placement_group: "PlacementGroup"): - ray.util.remove_placement_group(current_placement_group) - s = time.time() - wait_interval = 10 - while time.time() - s < PG_WAIT_TIMEOUT: - pg = ray.util.get_current_placement_group() - if pg is None: - break - - # Exponential backoff for warning print. - wait_interval *= 2 - logger.info( - "Waiting for removing a placement group of specs for %d seconds.", - int(time.time() - s), - ) - time.sleep(wait_interval) - - def initialize_ray_cluster( parallel_config: ParallelConfig, ray_address: str | None = None,