chore: import upstream snapshot with attribution

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

@@ -0,0 +1,54 @@
+from typing import Optional, Type
+
+import pytest
+
+from ray.air.execution._internal.barrier import Barrier
+
+
+def _raise(exception_type: Type[Exception] = RuntimeError, msg: Optional[str] = None):
+    def _raise_exception(*args, **kwargs):
+        raise exception_type(msg)
+
+    return _raise_exception
+
+
+def test_barrier_max_results():
+    """Test the `max_results` attribute.
+
+    - Set max_results=10
+    - Assert that the barrier completion callback is not invoked with num_results<10
+    - Assert that callback is invoked with num_results=10
+    - Assert that callback is not invoked again when more events arrive
+    - Assert that more events can arrive without triggering the callback after resetting
+    """
+    barrier = Barrier(max_results=10, on_completion=_raise(AssertionError))
+
+    for i in range(9):
+        barrier.arrive(i)
+
+    assert not barrier.completed
+
+    # Will trigger the on_completion callback
+    with pytest.raises(AssertionError):
+        barrier.arrive(10)
+
+    assert barrier.completed
+
+    assert barrier.num_results == 10
+
+    # Further events will not trigger callback again
+    barrier.arrive(11)
+
+    barrier.reset()
+
+    assert not barrier.completed
+
+    # After flushing more events can arrive
+    barrier.arrive(12)
+    assert barrier.num_results == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,266 @@
+import random
+from typing import Any, List, Optional
+
+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
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.exceptions import RayActorError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    """Simple actor for testing an execution flow.
+
+    This actor can fail in these ways:
+
+    1. On init if ``actor_init_kill`` is passed as a kwarg
+    2. On setup_1() if ``actor_setup_kill`` is passed as a kwarg (RayActorError)
+    3. On setup_1() if ``actor_setup_fail`` is passed as a kwarg (RayTaskError)
+    4. On train() if ``actor_train_kill`` is passed as a kwarg (RayTaskError)
+    5. On train() if ``actor_train_fail`` is passed as a kwarg (RayTaskError)
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+        if self.kwargs.get("actor_init_kill"):
+            raise RuntimeError("INIT")
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def setup_1(self):
+        if self.kwargs.get("actor_setup_kill"):
+            raise SystemExit
+
+        if self.kwargs.get("actor_setup_fail"):
+            raise RuntimeError("Setup")
+
+        return True
+
+    def setup_2(self):
+        return True
+
+    def train(self, value: float) -> float:
+        if value == 4:
+            if self.kwargs.get("actor_train_kill"):
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+
+            if self.kwargs.get("actor_train_fail"):
+                # RuntimeError will invoke a RayTaskError
+                raise RuntimeError("TASK")
+
+        return value
+
+
+class TrainFlow:
+    """This is a Ray Train-like execution flow.
+
+    - We want to run 4 actors in total ("trials")
+    - Each actor runs two init functions
+    - We train all actors in parallel for 10 iterations
+    - Errors can come up on actor construction, in the init functions,
+      or during training
+    - When an actor fails, restart that actor
+    - When a task fails, stop actor, and restart
+    """
+
+    def __init__(
+        self, actor_manager: RayActorManager, errors: Optional[List[str]] = None
+    ):
+        self._actor_manager = actor_manager
+        self._finished = False
+
+        self._actors_to_run = 4
+        self._tracked_actors = []
+        self._actors_stopped = 0
+
+        self._actors_to_replace = set()
+
+        self._ready_actors = set()
+        self._training_barrier = Barrier(
+            max_results=self._actors_to_run,
+            on_completion=self.training_barrier_completed,
+        )
+        self._restart_training = None
+
+        self._training_iter = 0
+        self._results = []
+
+        self._errors = errors
+
+    def setup_actors(self):
+        for actor_id in range(self._actors_to_run):
+            error_kwargs = {}
+            if self._errors:
+                error = random.choice(self._errors)
+                error_kwargs[error] = True
+
+            print("Actor", actor_id, "will be failing with", error_kwargs)
+
+            tracked_actor = self._actor_manager.add_actor(
+                cls=Actor,
+                kwargs={"id": actor_id, **error_kwargs},
+                resource_request=ResourceRequest([{"CPU": 1}]),
+                on_start=self.actor_started,
+                on_stop=self.actor_stopped,
+                on_error=self.actor_error,
+            )
+            self._tracked_actors.append(tracked_actor)
+
+    def actor_started(self, tracked_actor: TrackedActor):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "setup_1",
+            on_error=self.setup_error,
+            on_result=self.setup_1_result,
+        )
+
+    def actor_stopped(self, tracked_actor: TrackedActor):
+        self._ready_actors.discard(tracked_actor)
+
+        if tracked_actor in self._actors_to_replace:
+            self._replace_actor(tracked_actor=tracked_actor)
+        else:
+            self._actors_stopped += 1
+            self._finished = self._actors_stopped >= self._actors_to_run
+
+    def actor_error(self, tracked_actor: TrackedActor, exception: Exception):
+        self._ready_actors.discard(tracked_actor)
+        self._replace_actor(tracked_actor=tracked_actor)
+
+    def _replace_actor(self, tracked_actor: TrackedActor):
+        actor_index = self._tracked_actors.index(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"id": actor_index},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+
+        self._tracked_actors[actor_index] = replacement_actor
+
+    def setup_1_result(self, tracked_actor: TrackedActor, result: Any):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "setup_2",
+            on_error=self.setup_error,
+            on_result=self.setup_2_result,
+        )
+
+    def setup_2_result(self, tracked_actor: TrackedActor, result: Any):
+        self._ready_actors.add(tracked_actor)
+
+        if len(self._ready_actors) == self._actors_to_run:
+            self.continue_training()
+
+    def setup_error(self, tracked_actor: TrackedActor, exception: Exception):
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_to_replace.add(tracked_actor)
+        self._actor_manager.remove_actor(tracked_actor)
+
+    def continue_training(self):
+        if self._restart_training:
+            self._training_iter = self._restart_training
+        else:
+            self._training_iter += 1
+
+        self._training_barrier.reset()
+        self._actor_manager.schedule_actor_tasks(
+            self._tracked_actors,
+            "train",
+            args=(self._training_iter,),
+            on_result=self._training_barrier.arrive,
+            on_error=self.training_error,
+        )
+
+    def training_barrier_completed(self, barrier: Barrier):
+        self._results.append([res for _, res in barrier.get_results()])
+        self._restart_training = None
+
+        # If less than 10 epochs, continue training
+        if self._training_iter < 10:
+            return self.continue_training()
+
+        # Else, training finished
+        for tracked_actor in self._tracked_actors:
+            self._actor_manager.remove_actor(tracked_actor)
+
+    def training_error(self, tracked_actor: TrackedActor, exception: Exception):
+        self._restart_training = self._training_iter
+
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_to_replace.add(tracked_actor)
+        self._ready_actors.discard(tracked_actor)
+        self._actor_manager.remove_actor(tracked_actor)
+
+    def run(self):
+        self.setup_actors()
+
+        while not self._finished:
+            self._actor_manager.next()
+
+    def get_results(self) -> List[List[float]]:
+        return self._results
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "errors",
+    [
+        None,
+        "actor_init_kill",
+        "actor_setup_kill",
+        "actor_setup_fail",
+        "actor_train_kill",
+        "actor_train_fail",
+        # Chaos - every actor fails somehow, but in different ways
+        [
+            "actor_init_kill",
+            "actor_setup_kill",
+            "actor_setup_fail",
+            "actor_train_kill",
+            "actor_train_fail",
+        ],
+    ],
+)
+def test_e2e(ray_start_4_cpus, resource_manager_cls, errors):
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    if errors and isinstance(errors, str):
+        errors = [errors]
+
+    flow = TrainFlow(actor_manager=actor_manager, errors=errors)
+    flow.run()
+
+    results = flow.get_results()
+
+    assert results == [[i] * 4 for i in range(1, 11)], results
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,227 @@
+import random
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal.actor_manager import RayActorManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.exceptions import RayActorError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    """Simple actor for testing an execution flow.
+
+    This actor can fail in three ways:
+
+    1. On init if ``actor_error_init`` is passed as a kwarg
+    2. On run() if ``actor_error_task`` is passed as a kwarg (RayActorError)
+    3. On run() if ``task_error`` is passed as a kwarg (RayTaskError)
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+        if self.kwargs.get("actor_error_init"):
+            raise RuntimeError("INIT")
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def run(self, value: float) -> float:
+        if value == 2:
+            if self.kwargs.get("actor_error_task"):
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+
+            if self.kwargs.get("task_error"):
+                # RuntimeError will invoke a RayTaskError
+                raise RuntimeError("TASK")
+
+        return value
+
+
+class TuneFlow:
+    """This is a Ray Tune-like execution flow.
+
+    - We want to run 10 actors in total ("trials")
+    - Each actor collects 11 results sequentially
+    - We schedule up to 6 actors at the same time
+    - Every step, we see if we should add any new actors
+    - Otherwise, we just yield control to the event manager and process events one
+      by one
+    - When an actor is started, start training flow
+    - When a result comes in, schedule next future
+      - If this is the 11th result, stop actor
+    - When the last actor is stopped, set state to finished
+
+    - When an actor fails, restart
+    - When a task fails, stop actor, and restart
+    """
+
+    def __init__(
+        self, actor_manager: RayActorManager, errors: Optional[List[str]] = None
+    ):
+        self._actor_manager = actor_manager
+        self._finished = False
+
+        self._actors_to_run = 10
+        self._actors_started = 0
+        self._actors_stopped = 0
+        self._max_pending = 6
+
+        self._actor_to_id = {}
+        self._results = defaultdict(list)
+
+        self._errors = errors
+
+    def maybe_add_actors(self):
+        if self._actors_started >= self._actors_to_run:
+            return
+
+        if self._actor_manager.num_pending_actors >= self._max_pending:
+            return
+
+        error_kwargs = {}
+        if self._errors:
+            error = random.choice(self._errors)
+            error_kwargs[error] = True
+
+        actor_id = self._actors_started
+
+        print("Actor", actor_id, "will be failing with", error_kwargs)
+
+        tracked_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"id": actor_id, **error_kwargs},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+        self._actor_to_id[tracked_actor] = actor_id
+
+        self._actors_started += 1
+
+    def actor_started(self, tracked_actor: TrackedActor):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "run",
+            kwargs={"value": 0},
+            on_error=self.task_error,
+            on_result=self.task_result,
+        )
+
+    def actor_stopped(self, tracked_actor: TrackedActor):
+        self._actors_stopped += 1
+        self._finished = self._actors_stopped >= self._actors_to_run
+
+    def actor_error(self, tracked_actor: TrackedActor, exception: Exception):
+        actor_id = self._actor_to_id.pop(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={
+                "id": actor_id,
+                "actor_error_init": False,
+                "actor_error_task": False,
+                "task_error": False,
+            },
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+
+        self._actor_to_id[replacement_actor] = actor_id
+
+    def task_result(self, tracked_actor: TrackedActor, result: float):
+        actor_id = self._actor_to_id[tracked_actor]
+        self._results[actor_id].append(result)
+
+        if result == 10:
+            self._actor_manager.remove_actor(tracked_actor)
+        else:
+            self._actor_manager.schedule_actor_task(
+                tracked_actor,
+                "run",
+                kwargs={"value": result + 1},
+                on_result=self.task_result,
+                on_error=self.task_error,
+            )
+
+    def task_error(self, tracked_actor: TrackedActor, exception: Exception):
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_stopped -= 1  # account for extra stop
+        self._actor_manager.remove_actor(tracked_actor)
+        actor_id = self._actor_to_id.pop(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={
+                "id": actor_id,
+                "actor_error_init": False,
+                "actor_error_task": False,
+                "task_error": False,
+            },
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+        self._actor_to_id[replacement_actor] = actor_id
+
+    def run(self):
+        while not self._finished:
+            self.maybe_add_actors()
+            self._actor_manager.next(timeout=1)
+
+    def get_results(self) -> Dict[int, List[float]]:
+        return self._results
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "errors",
+    [
+        None,
+        "actor_error_init",
+        "actor_error_task",
+        "task_error",
+        # Chaos - every actor fails somehow, but in different ways
+        ["actor_error_init", "actor_error_task", "task_error"],
+    ],
+)
+def test_e2e(ray_start_4_cpus, resource_manager_cls, errors):
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    if errors and isinstance(errors, str):
+        errors = [errors]
+
+    flow = TuneFlow(actor_manager=actor_manager, errors=errors)
+    flow.run()
+
+    results = flow.get_results()
+
+    assert all(res[-1] == 10 for res in results.values()), results
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,224 @@
+import time
+from typing import Any, Type
+
+import pytest
+
+import ray
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.event_manager import RayEventManager
+from ray.exceptions import RayTaskError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+def succeeding(ret: Any = None) -> Any:
+    return ret
+
+
+@ray.remote
+def failing(exc: Type[Exception], *args) -> None:
+    raise exc(*args)
+
+
+@ray.remote
+def sleeping(seconds: int, result: Any) -> Any:
+    time.sleep(seconds)
+    return result
+
+
+def test_track_future_success(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that the on_result callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_future(succeeding.remote("a"), on_result=on_result)
+
+    event_manager.wait()
+    assert "a" in seen
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_success_no_callback(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that passing no callback still succeeds.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(succeeding.remote("a"))
+
+    event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that the on_error callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    class CustomError(RuntimeError):
+        pass
+
+    def on_error(exception: Exception):
+        seen.add(exception)
+
+    event_manager.track_future(failing.remote(CustomError), on_error=on_error)
+
+    event_manager.wait()
+    assert isinstance(seen.pop(), CustomError)
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error_no_callback(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that passing no callback raises the original error.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(failing.remote(RuntimeError))
+
+    with pytest.raises(RuntimeError):
+        event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+@pytest.mark.parametrize("results_per_wait", [None, 1, 5, 10, 100])
+def test_many_futures(ray_start_4_cpus, results_per_wait):
+    """Schedule 500 succeeding and failing futures.
+
+    Check that the callbacks get triggered correctly, independent of the number
+    of results we await per call to RayEventManager.wait().
+    """
+    num_futures = 500
+
+    event_manager = RayEventManager()
+
+    seen_results = set()
+    seen_errors = set()
+
+    def on_result(result: Any):
+        seen_results.add(result)
+
+    def on_error(exception: RayTaskError):
+        seen_errors.add(exception.cause.args[0])
+
+    for i in range(num_futures):
+        event_manager.track_futures(
+            [
+                succeeding.remote("a" + str(i)),
+                failing.remote(RuntimeError, "b" + str(i)),
+            ],
+            on_result=on_result,
+            on_error=on_error,
+        )
+
+    while event_manager.num_futures > 0:
+        event_manager.wait(num_results=results_per_wait)
+
+    for i in range(num_futures):
+        assert "a" + str(i) in seen_results
+        assert "b" + str(i) in seen_errors
+
+
+def test_timeout(ray_start_4_cpus):
+    """Test the timeout parameter.
+
+    Start 4 tasks: Two succeed immediately, two after 1 second.
+
+    After waiting for 0.5 seconds, the first two tasks should have returned.
+    After waiting for up to 5 seconds, the other two tasks should have returned.
+    But because the tasks take only 0.5 seconds to run, we should have waited
+    way less than 5 seconds.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            sleeping.remote(1, "c"),
+            sleeping.remote(1, "d"),
+        ],
+        on_result=on_result,
+    )
+
+    start = time.monotonic()
+    event_manager.wait(num_results=None, timeout=0.5)
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" not in seen
+    assert "d" not in seen
+
+    event_manager.wait(num_results=None, timeout=5)
+    taken = time.monotonic() - start
+
+    assert "c" in seen
+    assert "d" in seen
+
+    # Should have returned much earlier than after 5 seconds
+    assert taken < 3
+
+    assert not event_manager._tracked_futures
+
+
+def test_task_barrier(ray_start_4_cpus):
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_completion(barrier: Barrier):
+        seen.update(barrier.get_results())
+
+    barrier = Barrier(max_results=4, on_completion=on_completion)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            succeeding.remote("c"),
+            succeeding.remote("d"),
+            sleeping.remote(2, "e"),
+        ],
+        on_result=barrier.arrive,
+    )
+
+    event_manager.wait(num_results=4)
+
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" in seen
+    assert "d" in seen
+    assert "e" not in seen
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,178 @@
+import pytest
+
+import ray
+from ray.air.execution.resources.fixed import FixedResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_4_CPU = ResourceRequest([{"CPU": 4}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 2}])
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_acquire_return_resources(ray_start_4_cpus):
+    manager = FixedResourceManager(total_resources={"CPU": 4})
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_4_CPU)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+    ready_2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.free_resources(ready_2)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+
+def test_numerical_error(ray_start_4_cpus):
+    """Make sure we don't run into numerical errors when using fractional resources.
+
+    Legacy test: test_trial_runner::TrialRunnerTest::testResourceNumericalError
+    """
+    manager = FixedResourceManager(
+        total_resources={"CPU": 0.99, "GPU": 0.99, "a": 0.99}
+    )
+    resource_request = ResourceRequest([{"CPU": 0.33, "GPU": 0.33, "a": 0.33}])
+
+    for i in range(3):
+        manager.request_resources(resource_request)
+        assert manager.acquire_resources(
+            resource_request=resource_request
+        ), manager._available_resources
+
+    assert manager._available_resources["CPU"] == 0
+    assert manager._available_resources["GPU"] == 0
+    assert manager._available_resources["a"] == 0
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request resources with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    manager.request_resources(REQUEST_1_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request resources with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT placement strategies."""
+    manager = FixedResourceManager()
+
+    req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+    if strategy.startswith("STRICT_"):
+        with pytest.raises(RuntimeError):
+            manager.request_resources(req)
+    else:
+        manager.request_resources(req)
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy_nested(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT_SPREAD within a PG."""
+
+    @ray.remote
+    def nested_test():
+        manager = FixedResourceManager()
+
+        req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+        if strategy == "STRICT_SPREAD":
+            with pytest.raises(RuntimeError):
+                manager.request_resources(req)
+        else:
+            manager.request_resources(req)
+
+    pg = ray.util.placement_group([{"CPU": 2}])
+    ray.wait([pg.ready()])
+
+    try:
+        ray.get(
+            nested_test.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pg, placement_group_capture_child_tasks=True
+                )
+            ).remote()
+        )
+    finally:
+        ray.util.remove_placement_group(pg)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,409 @@
+import time
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.execution.resources.placement_group import PlacementGroupResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 2}])
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def _count_pg_states():
+    counter = Counter()
+    for _, pg_info in ray.util.placement_group_table().items():
+        counter[pg_info["state"]] += 1
+    return counter
+
+
+def test_request_cancel_resources(ray_start_4_cpus):
+    """Test that canceling a resource request clears the PG futures.
+
+    - Create request
+    - Assert actual PG is created
+    - Cancel request
+    - Assert staging future is removed
+    - Assert actual PG is removed
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Could be pending or created
+    pg_states = _count_pg_states()
+    assert pg_states["PENDING"] + pg_states["CREATED"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    assert manager.get_resource_futures()
+
+    manager.cancel_resource_request(REQUEST_2_CPU)
+
+    assert not manager.get_resource_futures()
+
+    pg_states = _count_pg_states()
+    assert pg_states["PENDING"] + pg_states["CREATED"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_acquire_return_resources(ray_start_4_cpus):
+    """Tests that acquiring and returning resources works.
+
+    - At the start, no resources should be ready (no PG scheduled)
+    - Request resources for 2 CPUs
+    - (wait until they are ready)
+    - Assert that these 2 CPUs are available to be acquired
+    - Acquire
+    - Assert that there are no 2 CPU resources available anymore
+    - Free resources
+    - Assert that the 2 CPU resources are still not available (no new request)
+        - This is also tested in includes test_request_cancel_resources
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # Request PG
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Wait until ready
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+
+    # PG exists
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    # Acquire PG
+    acquired = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # Free resources
+    manager.free_resources(acquired)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # PG still exists
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_request_pending(ray_start_4_cpus):
+    """Test that requesting too many resources leads to pending PGs.
+
+    - Cluster of 4 CPUs
+    - Request 3 PGs a 2 CPUs
+    - Acquire 2 PGs
+    - Assert no resources are available anymore
+    - Return both PGs
+    - Assert resources are available again
+    - Cancel request
+    - Assert no resources are available again
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Wait until some are ready
+    ray.wait(manager.get_resource_futures(), num_returns=2)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert len(manager.get_resource_futures()) == 1
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 2
+    assert pg_states["PENDING"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    acq1 = manager.acquire_resources(REQUEST_2_CPU)
+    acq2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.free_resources(acq1)
+    manager.free_resources(acq2)
+
+    # Third PG becomes ready
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 2
+
+    manager.cancel_resource_request(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 3
+
+
+def test_acquire_unavailable(ray_start_4_cpus):
+    """Test that acquiring resources that are not available returns None.
+
+    - Try to acquire
+    - Assert this does not work
+    - Request resources
+    - Wait until ready
+    - Acquire
+    - Assert this did work
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.acquire_resources(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+    assert manager.acquire_resources(REQUEST_2_CPU)
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request PG with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert res1 == {"CPU": 1}
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert res1 == {"CPU": 1}
+    assert res2 == {"CPU": 2}
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request PG with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert res1 == {}
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert res1 == {}
+    assert res2 == {"CPU": 2}
+
+
+def test_capture_child_tasks(ray_start_4_cpus):
+    """Test that child tasks are captured when creating placement groups.
+
+    - Request PG with 2 bundles (1 CPU and 2 CPUs)
+    - Bind a remote task that needs 2 CPUs to run
+    - Assert that it can be scheduled from within the first bundle
+
+    This is only the case if child tasks are captured in the placement groups, as
+    there is only 1 CPU available outside (on a 4 CPU cluster). The 2 CPUs
+    thus have to come from the placement group.
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def needs_cpus():
+        return "Ok"
+
+    @ray.remote
+    def spawn_child_task(num_cpus: int):
+        return ray.get(needs_cpus.options(num_cpus=num_cpus).remote())
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([spawn_child_task])
+
+    res = ray.get(av1.remote(2), timeout=2.0)
+
+    assert res
+
+
+def test_clear_state(ray_start_4_cpus):
+    """Test that clearing state will remove existing placement groups.
+
+    - Create resource request
+    - Wait until PG is scheduled
+    - Assert that Ray PG is created
+    - Call `mgr.clear()`
+    - Assert that resources are not ready anymore
+    - Assert that Ray PG is removed
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 0
+
+    manager.clear()
+
+    assert not manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_internal_state(ray_start_4_cpus):
+    """Test internal state mappings of the placement group manager.
+
+    This test makes assumptions and assertions around the internal state transition
+    of private properties of the placement group resource manager.
+
+    If you change internal handling logic of the manager, you may need to change this
+    test as well.
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+
+    assert manager.update_interval_s == 0
+
+    manager.has_resources_ready(REQUEST_2_CPU)
+
+    # The key may exist but the set should be empty
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+
+    ####
+    # 1. Request, wait until ready, cancel
+
+    # Request resources
+    manager.request_resources(REQUEST_2_CPU)
+
+    # PG should be staged
+    assert manager._request_to_staged_pgs[REQUEST_2_CPU]
+    pg = list(manager._request_to_staged_pgs[REQUEST_2_CPU])[0]
+    assert manager._pg_to_request[pg] == REQUEST_2_CPU
+
+    # Staging future should exist
+    assert manager._pg_to_staging_future[pg]
+    fut = manager._pg_to_staging_future[pg]
+    assert manager._staging_future_to_pg[fut] == pg
+
+    # Wait until PG is ready
+    while not manager.has_resources_ready(resource_request=REQUEST_2_CPU):
+        time.sleep(0.05)
+
+    # PG should now be ready
+    assert manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # PG should not be staged anymore
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    # Staging future should not exist anymore
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+
+    # Cancel request
+    manager.cancel_resource_request(REQUEST_2_CPU)
+
+    # PG should not be ready anymore
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # All PGs should be fully removed
+    assert not manager._pg_to_request
+
+    ####
+    # 2. Request, cancel while staging
+
+    # Stage another PG
+    manager.request_resources(REQUEST_2_CPU)
+    # Cancel request before it's ready
+    manager.cancel_resource_request(REQUEST_2_CPU)
+    # Assert no leftover
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    assert not manager._pg_to_request
+
+    ####
+    # 2. Request, acquire, free
+
+    # Stage another PG
+    manager.request_resources(REQUEST_2_CPU)
+    pg = list(manager._request_to_staged_pgs[REQUEST_2_CPU])[0]
+    # Wait until PG is ready
+    while not manager.has_resources_ready(resource_request=REQUEST_2_CPU):
+        time.sleep(0.05)
+    # Acquire
+    acquired_resources = manager.acquire_resources(resource_request=REQUEST_2_CPU)
+    # Assert no staging/ready leftover
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # We still retain this mapping
+    assert manager._pg_to_request
+    # And we keep track of acquired PGs
+    assert pg in manager._acquired_pgs
+
+    # Free PG
+    manager.free_resources(acquired_resources)
+    # State should be cleared now
+    assert not manager._pg_to_request
+    assert not manager._acquired_pgs
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,38 @@
+import pytest
+
+from ray.air.execution.resources.request import ResourceRequest
+
+
+def test_request_same():
+    """Test that resource requests are the same if they share the same properties."""
+
+    assert ResourceRequest([{"CPU": 1}]) == ResourceRequest([{"CPU": 1}])
+
+    # multiple bundles work
+    assert ResourceRequest([{"CPU": 1}, {"CPU": 2}]) == ResourceRequest(
+        [{"CPU": 1}, {"CPU": 2}]
+    )
+
+    # multiple resources work
+    assert ResourceRequest([{"CPU": 1, "GPU": 1}]) == ResourceRequest(
+        [{"CPU": 1, "GPU": 1}]
+    )
+
+    # 0 resources are ignored
+    assert ResourceRequest([{"CPU": 0, "GPU": 1}]) == ResourceRequest([{"GPU": 1}])
+
+    # PACK is implicit
+    assert ResourceRequest([{"CPU": 1}], strategy="PACK") == ResourceRequest(
+        [{"CPU": 1}]
+    )
+
+    # Non match: different strategy
+    assert ResourceRequest([{"CPU": 1}], strategy="PACK") != ResourceRequest(
+        [{"CPU": 1}], strategy="SPREAD"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

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

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

@@ -0,0 +1,122 @@
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal.actor_manager import RayActorManager
+
+RESOURCE_MANAGERS = [FixedResourceManager, PlacementGroupResourceManager]
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    def foo(self, val, error: bool = False):
+        if error:
+            raise RuntimeError
+        return val
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_resolve(ray_start_4_cpus, resource_manager_cls):
+    """Test that the `on_result` callback is invoked when a task completes.
+
+    - Instantiate global data object
+    - Schedule task that returns a value
+    - The callback writes the returned value to the global data object
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    seen = {"data": 0}
+
+    def result_callback(tracked_actor, result):
+        seen["data"] = result
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(
+        tracked_actor, "foo", (4, False), on_result=result_callback
+    )
+    actor_manager.next()
+    actor_manager.next()
+
+    assert seen["data"] == 4
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+@pytest.mark.parametrize("num_tasks", [1, 10, 100])
+def test_resolve_many(ray_start_4_cpus, resource_manager_cls, num_tasks):
+    """Schedule ``num_tasks`` tasks and wait until ``wait_for_events`` of them resolve.
+
+    Every resolved task will increase a counter by its return value (1).
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    seen = {"data": 0}
+
+    def result_callback(tracked_actor, result):
+        seen["data"] += result
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.next()
+
+    for i in range(num_tasks):
+        actor_manager.schedule_actor_task(
+            tracked_actor, "foo", (1, False), on_result=result_callback
+        )
+
+    for i in range(num_tasks):
+        actor_manager.next()
+        assert seen["data"] == i + 1
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_error_noop(ray_start_4_cpus, resource_manager_cls):
+    """When no `on_error` callback is specified, errors should be ignored."""
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(tracked_actor, "foo", (1, True))
+    actor_manager.next()
+    actor_manager.next()
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_error_custom(ray_start_4_cpus, resource_manager_cls):
+    """When an `on_error` callback is specified, it is invoked."""
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    stats = Counter()
+
+    def error_callback(tracked_actor, exception):
+        stats["exception"] += 1
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(
+        tracked_actor, "foo", (1, True), on_error=error_callback
+    )
+
+    actor_manager.next()
+    actor_manager.next()
+    assert stats["exception"] == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))