wehub-resource-sync
333 lines
9.4 KiB
Python
333 lines
9.4 KiB
Python
import threading
|
|
import time
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.exceptions import RayActorError
|
|
from ray.tests.client_test_utils import create_remote_signal_actor
|
|
from ray.train.v2._internal.execution.controller.placement_group_cleaner import (
|
|
PlacementGroupCleaner,
|
|
)
|
|
from ray.util.placement_group import placement_group, remove_placement_group
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def ray_start():
|
|
ray.init(num_cpus=4)
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def monitoring_started_signal(monkeypatch):
|
|
"""Signal actor set whenever PlacementGroupCleaner begins monitoring."""
|
|
|
|
SignalActor = create_remote_signal_actor(ray)
|
|
signal_actor = SignalActor.options(num_cpus=0).remote()
|
|
|
|
original_start_monitoring = PlacementGroupCleaner.start_monitoring
|
|
|
|
def start_monitoring_with_signal(self):
|
|
started = original_start_monitoring(self)
|
|
if started:
|
|
signal_actor.send.remote()
|
|
return started
|
|
|
|
monkeypatch.setattr(
|
|
PlacementGroupCleaner,
|
|
"start_monitoring",
|
|
start_monitoring_with_signal,
|
|
)
|
|
|
|
try:
|
|
yield signal_actor
|
|
finally:
|
|
try:
|
|
ray.kill(signal_actor)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class MockController:
|
|
"""Mock controller actor for testing."""
|
|
|
|
def __init__(self):
|
|
self._alive = True
|
|
|
|
def get_actor_id(self):
|
|
return ray.get_runtime_context().get_actor_id()
|
|
|
|
def shutdown(self):
|
|
"""Simulate controller death."""
|
|
self._alive = False
|
|
ray.actor.exit_actor()
|
|
|
|
|
|
def test_placement_group_cleaner_basic_lifecycle(monitoring_started_signal):
|
|
"""Test that the PlacementGroupCleaner can be launched and stopped."""
|
|
|
|
# Create a mock controller
|
|
controller = MockController.remote()
|
|
controller_id = ray.get(controller.get_actor_id.remote())
|
|
|
|
# Launch cleaner as detached
|
|
cleaner = (
|
|
ray.remote(num_cpus=0)(PlacementGroupCleaner)
|
|
.options(
|
|
name="test_pg_cleaner",
|
|
namespace="test",
|
|
lifetime="detached",
|
|
get_if_exists=False,
|
|
)
|
|
.remote(
|
|
controller_actor_id=controller_id,
|
|
check_interval_s=0.1,
|
|
get_actor_timeout_s=2,
|
|
stop_timeout=5,
|
|
)
|
|
)
|
|
|
|
# Create a placement group
|
|
pg = placement_group([{"CPU": 1}], strategy="SPREAD")
|
|
ray.get(pg.ready())
|
|
|
|
# Register placement group for cleanup
|
|
ray.get(cleaner.register_placement_group.remote(pg))
|
|
|
|
# Start monitoring asynchronously (same behavior as production code)
|
|
cleaner.start_monitoring.remote()
|
|
|
|
# Wait deterministically for monitoring to start to avoid flakiness.
|
|
ray.get(monitoring_started_signal.wait.remote())
|
|
|
|
# Controller is still alive, so PG should still exist
|
|
assert pg.id is not None
|
|
|
|
# Stop the cleaner gracefully; tolerate the actor exiting itself.
|
|
try:
|
|
ray.get(cleaner.stop.remote(), timeout=2.0)
|
|
except RayActorError:
|
|
pass
|
|
|
|
# PG should still exist after graceful stop
|
|
try:
|
|
# If PG exists, this should work
|
|
remove_placement_group(pg)
|
|
except Exception as e:
|
|
pytest.fail(f"Placement group should still exist after graceful stop: {e}")
|
|
finally:
|
|
try:
|
|
ray.get(controller.shutdown.remote())
|
|
except RayActorError:
|
|
pass
|
|
|
|
|
|
def test_pg_cleaner_cleans_up_on_controller_death(monitoring_started_signal):
|
|
"""Test that the PG cleaner removes PG when controller dies."""
|
|
|
|
# Create a mock controller
|
|
controller = MockController.remote()
|
|
controller_id = ray.get(controller.get_actor_id.remote())
|
|
|
|
# Launch cleaner as detached
|
|
cleaner = (
|
|
ray.remote(num_cpus=0)(PlacementGroupCleaner)
|
|
.options(
|
|
name="test_pg_cleaner_cleanup",
|
|
namespace="test",
|
|
lifetime="detached",
|
|
get_if_exists=False,
|
|
)
|
|
.remote(
|
|
controller_actor_id=controller_id,
|
|
check_interval_s=0.1,
|
|
get_actor_timeout_s=2,
|
|
stop_timeout=5,
|
|
)
|
|
)
|
|
|
|
# Create a placement group
|
|
pg = placement_group([{"CPU": 1}], strategy="SPREAD")
|
|
ray.get(pg.ready())
|
|
|
|
# Register placement group for cleanup
|
|
ray.get(cleaner.register_placement_group.remote(pg))
|
|
|
|
cleaner.start_monitoring.remote()
|
|
|
|
# Wait deterministically for monitoring to start to avoid flakiness.
|
|
ray.get(monitoring_started_signal.wait.remote())
|
|
|
|
# Kill the controller
|
|
ray.kill(controller)
|
|
|
|
# Wait for cleaner to detect death and clean up
|
|
time.sleep(1.0)
|
|
|
|
# Try to verify PG was cleaned up
|
|
# Note: After cleanup, the PG ID should no longer be valid
|
|
try:
|
|
# Try to remove the PG - if it was already removed, this should fail
|
|
remove_placement_group(pg)
|
|
except Exception:
|
|
# This is expected - PG was already cleaned up
|
|
pass
|
|
|
|
# The cleaner should exit after performing cleanup since it's detached.
|
|
with pytest.raises(RayActorError):
|
|
ray.get(cleaner.start_monitoring.remote(), timeout=2.0)
|
|
|
|
|
|
def test_pg_cleaner_exits_on_controller_death_without_pg_registration(
|
|
monitoring_started_signal,
|
|
):
|
|
"""Cleaner should exit cleanly if controller dies before any PG is registered."""
|
|
|
|
controller = MockController.remote()
|
|
controller_id = ray.get(controller.get_actor_id.remote())
|
|
|
|
cleaner = (
|
|
ray.remote(num_cpus=0)(PlacementGroupCleaner)
|
|
.options(
|
|
name="test_pg_cleaner_no_pg",
|
|
namespace="test",
|
|
lifetime="detached",
|
|
get_if_exists=False,
|
|
)
|
|
.remote(
|
|
controller_actor_id=controller_id,
|
|
check_interval_s=0.1,
|
|
get_actor_timeout_s=2,
|
|
stop_timeout=5,
|
|
)
|
|
)
|
|
|
|
cleaner.start_monitoring.remote()
|
|
ray.get(monitoring_started_signal.wait.remote())
|
|
|
|
ray.kill(controller)
|
|
|
|
# Give the cleaner a moment to observe controller death and self-exit.
|
|
time.sleep(1.0)
|
|
|
|
with pytest.raises(RayActorError):
|
|
ray.get(cleaner.start_monitoring.remote(), timeout=2.0)
|
|
|
|
|
|
def test_pg_cleaner_handles_duplicate_start():
|
|
"""Test that cleaner handles duplicate start_monitoring calls."""
|
|
|
|
controller = MockController.remote()
|
|
controller_id = ray.get(controller.get_actor_id.remote())
|
|
|
|
cleaner = (
|
|
ray.remote(num_cpus=0)(PlacementGroupCleaner)
|
|
.options(
|
|
name="test_pg_cleaner_duplicate",
|
|
namespace="test",
|
|
lifetime="detached",
|
|
get_if_exists=False,
|
|
)
|
|
.remote(
|
|
controller_actor_id=controller_id,
|
|
check_interval_s=0.1,
|
|
get_actor_timeout_s=2,
|
|
stop_timeout=5,
|
|
)
|
|
)
|
|
|
|
pg = placement_group([{"CPU": 1}], strategy="SPREAD")
|
|
ray.get(pg.ready())
|
|
ray.get(cleaner.register_placement_group.remote(pg))
|
|
|
|
# Start monitoring asynchronously
|
|
cleaner.start_monitoring.remote()
|
|
|
|
# Stop
|
|
try:
|
|
ray.get(cleaner.stop.remote(), timeout=2.0)
|
|
except RayActorError:
|
|
pass
|
|
|
|
# Detached cleaner should be gone after stop.
|
|
with pytest.raises(RayActorError):
|
|
ray.get(cleaner.start_monitoring.remote(), timeout=2.0)
|
|
|
|
|
|
def test_pg_cleaner_race_condition_new_pg_not_orphaned():
|
|
"""PG registered during the race window must not be orphaned.
|
|
|
|
Race window:
|
|
1. PGC calls queue.get()
|
|
2. controller puts new_pg into queue
|
|
3. controller dies
|
|
4. PGC checks is_actor_alive
|
|
|
|
The new_pg never gets cleaned up.
|
|
"""
|
|
old_pg = placement_group([{"CPU": 1}])
|
|
new_pg = placement_group([{"CPU": 1}])
|
|
ray.get(old_pg.ready())
|
|
ray.get(new_pg.ready())
|
|
|
|
in_race_window = threading.Event()
|
|
release = threading.Event()
|
|
call_count = {"n": 0}
|
|
|
|
def controlled_is_actor_alive(*args, **kwargs):
|
|
call_count["n"] += 1
|
|
if call_count["n"] == 1:
|
|
return True # first iteration: controller alive, keep looping
|
|
# Second call: queue.get() has already timed out
|
|
in_race_window.set()
|
|
release.wait(timeout=5)
|
|
return False # simulate controller death
|
|
|
|
target = (
|
|
"ray.train.v2._internal.execution.controller"
|
|
".placement_group_cleaner.is_actor_alive"
|
|
)
|
|
with patch(target, side_effect=controlled_is_actor_alive):
|
|
# PlacementGroupCleaner is instantiated directly here (not as a Ray actor),
|
|
# so ray.actor.exit_actor() is not tested here
|
|
cleaner = PlacementGroupCleaner(
|
|
controller_actor_id="test_pg_cleaner",
|
|
check_interval_s=0.05,
|
|
get_actor_timeout_s=2,
|
|
stop_timeout=5,
|
|
)
|
|
|
|
cleaner.register_placement_group(old_pg)
|
|
cleaner.start_monitoring()
|
|
|
|
assert in_race_window.wait(timeout=5), "Timed out waiting for race window"
|
|
|
|
# new_pg is now in the queue, but cleaner has already passed queue.get()
|
|
cleaner.register_placement_group(new_pg)
|
|
|
|
monitor_thread = cleaner._monitor_thread
|
|
assert monitor_thread is not None
|
|
release.set()
|
|
|
|
monitor_thread.join(timeout=5)
|
|
assert not monitor_thread.is_alive()
|
|
|
|
new_pg_state = ray.util.placement_group_table(new_pg).get("state")
|
|
assert (
|
|
new_pg_state == "REMOVED"
|
|
), f"Race condition: new_pg was orphaned (state={new_pg_state})"
|
|
|
|
try:
|
|
remove_placement_group(old_pg)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|