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