chore: import upstream snapshot with attribution

python/ray/air/tests/execution/test_tracked_actor.py

@@ -0,0 +1,366 @@
+import gc
+import threading
+import time
+from collections import Counter
+from typing import Any, Optional, Type
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.actor_manager import RayActorManager
+
+
+def _raise(exception_type: Type[Exception] = RuntimeError, msg: Optional[str] = None):
+    def _raise_exception(*args, **kwargs):
+        raise exception_type(msg)
+
+    return _raise_exception
+
+
+class Started(RuntimeError):
+    pass
+
+
+class Stopped(RuntimeError):
+    pass
+
+
+class Failed(RuntimeError):
+    pass
+
+
+class Result(RuntimeError):
+    pass
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def cleanup():
+    # Garbage collect at the start
+    # This ensures that all resources are freed up for the upcoming test.
+    gc.collect()
+    yield
+
+
+class Actor:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def task(self, value: Any):
+        return value
+
+
+@ray.remote(num_cpus=4)
+def fn():
+    return True
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("actor_cls", [Actor, ray.remote(Actor)])
+@pytest.mark.parametrize("kill", [False, True])
+def test_start_stop_actor(ray_start_4_cpus, resource_manager_cls, actor_cls, kill):
+    """Test that starting and stopping actors work and invokes a callback.
+
+    - Start an actor
+    - Starting should trigger start callback
+    - Schedule actor task, which should resolve (meaning actor successfully started)
+    - Stop actor, which should resolve and trigger stop callback
+    - Schedule remote fn that takes up all cluster resources. This should resolve,
+      meaning that the actor was stopped successfully.
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    # Start actor, set callbacks
+    tracked_actor = actor_manager.add_actor(
+        cls=actor_cls,
+        kwargs={"key": "val"},
+        resource_request=ResourceRequest([{"CPU": 4}]),
+        on_start=_raise(Started),
+        on_stop=_raise(Stopped),
+        on_error=_raise(Failed),
+    )
+
+    # Actor should be started
+    with pytest.raises(Started):
+        actor_manager.next()
+
+    # Schedule task on actor which should resolve (actor successfully started)
+    actor_manager.schedule_actor_task(
+        tracked_actor, "task", (1,), on_result=_raise(Result)
+    )
+
+    with pytest.raises(Result):
+        actor_manager.next()
+
+    # Now we can assert that there are no CPUS resources available anymore.
+    # Note that actor starting is asynchronous, so we can't assert this right away
+    # - that's why we wait for the actor task to resolve first.
+    assert ray.available_resources().get("CPU", 0.0) == 0, ray.available_resources()
+
+    # Stop actor
+    actor_manager.remove_actor(tracked_actor, kill=kill)
+
+    with pytest.raises(Stopped):
+        actor_manager.next()
+
+    # This task takes up all the cluster resources. It should resolve now that
+    # the actor was terminated.
+    assert ray.get(fn.remote(), timeout=5)
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_start_many_actors(ray_start_4_cpus, resource_manager_cls):
+    """Test that starting more actors than fit onto the cluster works.
+
+    - Request 10 actors
+    - 4 can be started. Assert they are started
+    - Stop 2
+    - Assert 2 are stopped and 2 new ones are started
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    running_actors = []
+    # stats keeps track of started/stopped actors
+    stats = Counter()
+
+    def start_callback(tracked_actor):
+        running_actors.append(tracked_actor)
+        stats["started"] += 1
+
+    def stop_callback(tracked_actor):
+        running_actors.remove(tracked_actor)
+        stats["stopped"] += 1
+
+    # start 10 actors
+    expected_actors = []
+    for i in range(10):
+        tracked_actor = actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"key": "val"},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=start_callback,
+            on_stop=stop_callback,
+            on_error=_raise(Failed),
+        )
+        expected_actors.append(tracked_actor)
+
+    # wait for some actor starts
+    for i in range(4):
+        actor_manager.next()
+
+    # we should now have 4 started actors
+    assert stats["started"] == 4
+    assert stats["stopped"] == 0
+    assert len(running_actors) == 4
+    assert set(running_actors) == set(expected_actors[:4])
+
+    # stop 2 actors
+    actor_manager.remove_actor(running_actors[0])
+    actor_manager.remove_actor(running_actors[1])
+
+    # Wait four times, twice for termination, twice for start
+    for i in range(4):
+        actor_manager.next()
+
+    # we should have 4 running actors, 6 started and 2 stopped
+    assert stats["started"] == 6
+    assert stats["stopped"] == 2
+    assert len(running_actors) == 4
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("where", ["init", "fn"])
+def test_actor_fail(ray_start_4_cpus, cleanup, resource_manager_cls, where):
+    """Test that actor failures are handled properly.
+
+    - Start actor that either fails on init or in a task (RayActorError)
+    - Schedule task on actor
+    - Assert that the correct callbacks are called
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    # keep track of failed tasks and actors
+    stats = Counter()
+
+    @ray.remote
+    class FailingActor:
+        def __init__(self, where):
+            self._where = where
+            if self._where == "init":
+                raise RuntimeError("INIT")
+
+        def fn(self):
+            if self._where == "fn":
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+            return True
+
+    def fail_callback_actor(tracked_actor, exception):
+        stats["failed_actor"] += 1
+
+    def fail_callback_task(tracked_actor, exception):
+        stats["failed_task"] += 1
+
+    # Start actor
+    tracked_actor = actor_manager.add_actor(
+        cls=FailingActor,
+        kwargs={"where": where},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_error=fail_callback_actor,
+    )
+
+    if where != "init":
+        # Wait until it is started. This won't invoke any callback, yet
+        actor_manager.next()
+
+        assert stats["failed_actor"] == 0
+        assert stats["failed_task"] == 0
+
+        # Schedule task
+        actor_manager.schedule_actor_task(
+            tracked_actor, "fn", on_error=fail_callback_task
+        )
+
+    # Yield control and wait for task resolution. This will invoke the callback.
+    actor_manager.next()
+
+    assert stats["failed_actor"] == 1
+    assert stats["failed_task"] == bool(where != "init")
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_stop_actor_before_start(
+    ray_start_4_cpus, tmp_path, cleanup, resource_manager_cls
+):
+    """Test that actor failures are handled properly.
+
+    - Start actor that either fails on init or in a task (RayActorError)
+    - Schedule task on actor
+    - Assert that the correct callbacks are called
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    hang_marker = tmp_path / "hang.txt"
+
+    @ray.remote
+    class HangingActor:
+        def __init__(self):
+            while not hang_marker.exists():
+                time.sleep(0.05)
+
+    tracked_actor = actor_manager.add_actor(
+        HangingActor,
+        kwargs={},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_start=_raise(RuntimeError, "Should not have started"),
+        on_stop=_raise(RuntimeError, "Should not have stopped"),
+    )
+    while not actor_manager.is_actor_started(tracked_actor):
+        actor_manager.next(0.05)
+
+    # Actor started but hasn't triggered on_start, yet
+    actor_manager.remove_actor(tracked_actor)
+    hang_marker.write_text("")
+    while actor_manager.is_actor_started(tracked_actor):
+        actor_manager.next(0.05)
+
+    assert actor_manager.num_live_actors == 0
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("start_thread", [False, True])
+def test_stop_actor_custom_future(
+    ray_start_4_cpus, tmp_path, cleanup, resource_manager_cls, start_thread
+):
+    """If we pass a custom stop future, the actor should still be shutdown by GC.
+
+    This should also be the case when we start a thread in the background, as we
+    do e.g. in Ray Tune's function runner.
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    hang_marker = tmp_path / "hang.txt"
+
+    actor_name = f"stopping_actor_{resource_manager_cls.__name__}_{start_thread}"
+
+    @ray.remote(name=actor_name)
+    class HangingStopActor:
+        def __init__(self):
+            self._thread = None
+            self._stop_event = threading.Event()
+            if start_thread:
+
+                def entrypoint():
+                    while True:
+                        print("Thread!")
+                        time.sleep(1)
+                        if self._stop_event.is_set():
+                            sys.exit(0)
+
+                self._thread = threading.Thread(target=entrypoint)
+                self._thread.start()
+
+        def stop(self):
+            print("Waiting")
+            while not hang_marker.exists():
+                time.sleep(0.05)
+            self._stop_event.set()
+            print("stopped")
+
+    start_barrier = Barrier(max_results=1)
+    stop_barrier = Barrier(max_results=1)
+
+    tracked_actor = actor_manager.add_actor(
+        HangingStopActor,
+        kwargs={},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_start=start_barrier.arrive,
+        on_stop=stop_barrier.arrive,
+    )
+    while not start_barrier.completed:
+        actor_manager.next(0.05)
+
+    # Actor is alive
+    assert ray.get_actor(actor_name)
+
+    stop_future = actor_manager.schedule_actor_task(tracked_actor, "stop")
+    actor_manager.remove_actor(tracked_actor, kill=False, stop_future=stop_future)
+
+    assert not stop_barrier.completed
+
+    hang_marker.write_text("!")
+
+    while not stop_barrier.completed:
+        actor_manager.next(0.05)
+
+    # Actor should have stopped now and should get cleaned up
+    with pytest.raises(ValueError):
+        ray.get_actor(actor_name)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))