chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_placement_group_cleaner.py

@@ -0,0 +1,332 @@
+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__]))