[CI][Ray] Fix flaky multi-node assignment test after placement-group teardown (#53621)

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))

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,