wehub-resource-sync
228 lines
6.8 KiB
Python
228 lines
6.8 KiB
Python
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__]))
|