[CI][Ray] Fix flaky multi-node assignment test after placement-group teardown (#53621)
yzong-rh committed
Aug 29, 2026 at 19:48 UTC
680e2177e473ed8dfaa9773f7ead185b369cab46
2 files changed
+17
-20
tests/distributed/test_multi_node_assignment.py
+17
-2
@@ -11,6 +11,7 @@ pytest distributed/test_multi_node_assignment.py
11
"""
12
13
import os
14
+import time
15
16
import pytest
17
import ray
@@ -19,7 +20,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
20
from vllm import initialize_ray_cluster
21
from vllm.config import ParallelConfig
22
from vllm.utils.network_utils import get_ip
22
-from vllm.v1.executor.ray_utils import _wait_until_pg_removed
23
+from vllm.v1.executor.ray_utils import available_resources_per_node
24
25
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
26
@@ -61,4 +62,18 @@ def test_multi_node_assignment() -> None:
62
for worker in workers:
63
ray.kill(worker)
64
64
- _wait_until_pg_removed(config.placement_group)
65
+ ray.util.remove_placement_group(config.placement_group)
66
+
67
+ gpus_needed = len(workers)
68
+ deadline = time.monotonic() + 60
69
+ node_id = ray.get_runtime_context().get_node_id()
70
+ while True:
71
+ free_gpus = available_resources_per_node().get(node_id, {}).get("GPU", 0)
72
+ if free_gpus >= gpus_needed:
73
+ break
74
+
75
+ remaining = deadline - time.monotonic()
76
+ if remaining <= 0:
77
+ raise TimeoutError("placement group resources were not released")
78
+
79
+ time.sleep(min(1, remaining))
vllm/v1/executor/ray_utils.py
-18
@@ -517,24 +517,6 @@ def _wait_until_pg_ready(current_placement_group: "PlacementGroup"):
517
) from None
518
519
520
-def _wait_until_pg_removed(current_placement_group: "PlacementGroup"):
521
- ray.util.remove_placement_group(current_placement_group)
522
- s = time.time()
523
- wait_interval = 10
524
- while time.time() - s < PG_WAIT_TIMEOUT:
525
- pg = ray.util.get_current_placement_group()
526
- if pg is None:
527
- break
528
-
529
- # Exponential backoff for warning print.
530
- wait_interval *= 2
531
- logger.info(
532
- "Waiting for removing a placement group of specs for %d seconds.",
533
- int(time.time() - s),
534
- )
535
- time.sleep(wait_interval)
536
-
537
-
520
def initialize_ray_cluster(
521
parallel_config: ParallelConfig,
522
ray_address: str | None = None,