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