chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,12 @@
|
||||
from ray.air.execution.resources.fixed import FixedResourceManager
|
||||
from ray.air.execution.resources.placement_group import PlacementGroupResourceManager
|
||||
from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
|
||||
from ray.air.execution.resources.resource_manager import ResourceManager
|
||||
|
||||
__all__ = [
|
||||
"ResourceRequest",
|
||||
"AcquiredResources",
|
||||
"ResourceManager",
|
||||
"FixedResourceManager",
|
||||
"PlacementGroupResourceManager",
|
||||
]
|
||||
@@ -0,0 +1,5 @@
|
||||
from ray.air.execution._internal.actor_manager import RayActorManager
|
||||
from ray.air.execution._internal.barrier import Barrier
|
||||
from ray.air.execution._internal.tracked_actor import TrackedActor
|
||||
|
||||
__all__ = ["Barrier", "RayActorManager", "TrackedActor"]
|
||||
@@ -0,0 +1,906 @@
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from collections import Counter, defaultdict
|
||||
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
|
||||
|
||||
import ray
|
||||
from ray.air.execution._internal.event_manager import RayEventManager
|
||||
from ray.air.execution._internal.tracked_actor import TrackedActor
|
||||
from ray.air.execution._internal.tracked_actor_task import TrackedActorTask
|
||||
from ray.air.execution.resources import (
|
||||
AcquiredResources,
|
||||
ResourceManager,
|
||||
ResourceRequest,
|
||||
)
|
||||
from ray.exceptions import RayActorError, RayTaskError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RayActorManager:
|
||||
"""Management class for Ray actors and actor tasks.
|
||||
|
||||
This class provides an event-based management interface for actors, and
|
||||
actor tasks.
|
||||
|
||||
The manager can be used to start actors, stop actors, and schedule and
|
||||
track task futures on these actors.
|
||||
The manager will then invoke callbacks related to the tracked entities.
|
||||
|
||||
For instance, when an actor is added with
|
||||
:meth:`add_actor() <RayActorManager.add_actor>`,
|
||||
a :ref:`TrackedActor <ray.air.execution._internal.tracked_actor.TrackedActor`
|
||||
object is returned. An ``on_start`` callback can be specified that is invoked
|
||||
once the actor successfully started. Similarly, ``on_stop`` and ``on_error``
|
||||
can be used to specify callbacks relating to the graceful or ungraceful
|
||||
end of an actor's lifetime.
|
||||
|
||||
When scheduling an actor task using
|
||||
:meth:`schedule_actor_task()
|
||||
<ray.air.execution._internal.actor_manager.RayActorManager.schedule_actor_task>`,
|
||||
an ``on_result`` callback can be specified that is invoked when the task
|
||||
successfully resolves, and an ``on_error`` callback will resolve when the
|
||||
task fails.
|
||||
|
||||
The RayActorManager does not implement any true asynchronous processing. Control
|
||||
has to be explicitly yielded to the event manager via :meth:`RayActorManager.next`.
|
||||
Callbacks will only be invoked when control is with the RayActorManager, and
|
||||
callbacks will always be executed sequentially in order of arriving events.
|
||||
|
||||
Args:
|
||||
resource_manager: Resource manager used to request resources for the actors.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from ray.air.execution import ResourceRequest
|
||||
from ray.air.execution._internal import RayActorManager
|
||||
|
||||
actor_manager = RayActorManager()
|
||||
|
||||
# Request an actor
|
||||
tracked_actor = actor_manager.add_actor(
|
||||
ActorClass,
|
||||
kwargs={},
|
||||
resource_request=ResourceRequest([{"CPU": 1}]),
|
||||
on_start=actor_start_callback,
|
||||
on_stop=actor_stop_callback,
|
||||
on_error=actor_error_callback
|
||||
)
|
||||
|
||||
# Yield control to event manager to start actor
|
||||
actor_manager.next()
|
||||
|
||||
# Start task on the actor (ActorClass.foo.remote())
|
||||
tracked_actor_task = actor_manager.schedule_actor_task(
|
||||
tracked_actor,
|
||||
method_name="foo",
|
||||
on_result=task_result_callback,
|
||||
on_error=task_error_callback
|
||||
)
|
||||
|
||||
# Again yield control to event manager to process task futures
|
||||
actor_manager.wait()
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, resource_manager: ResourceManager):
|
||||
self._resource_manager: ResourceManager = resource_manager
|
||||
|
||||
self._actor_state_events = RayEventManager()
|
||||
self._actor_task_events = RayEventManager()
|
||||
|
||||
# ---
|
||||
# Tracked actor futures.
|
||||
|
||||
# This maps TrackedActor objects to their futures. We use this to see if an
|
||||
# actor has any futures scheduled and to remove them when we terminate an actor.
|
||||
|
||||
# Actors to actor task futures
|
||||
self._tracked_actors_to_task_futures: Dict[
|
||||
TrackedActor, Set[ray.ObjectRef]
|
||||
] = defaultdict(set)
|
||||
|
||||
# Actors to actor state futures (start/terminate)
|
||||
self._tracked_actors_to_state_futures: Dict[
|
||||
TrackedActor, Set[ray.ObjectRef]
|
||||
] = defaultdict(set)
|
||||
|
||||
# ---
|
||||
# Pending actors.
|
||||
# We use three dicts for actors that are requested but not yet started.
|
||||
|
||||
# This dict keeps a list of actors associated with each resource request.
|
||||
# We use this to start actors in the correct order when their resources
|
||||
# become available.
|
||||
self._resource_request_to_pending_actors: Dict[
|
||||
ResourceRequest, List[TrackedActor]
|
||||
] = defaultdict(list)
|
||||
|
||||
# This dict stores the actor class, kwargs, and resource request of
|
||||
# pending actors. Once the resources are available, we start the remote
|
||||
# actor class with its args. We need the resource request to cancel it
|
||||
# if needed.
|
||||
self._pending_actors_to_attrs: Dict[
|
||||
TrackedActor, Tuple[Type, Dict[str, Any], ResourceRequest]
|
||||
] = {}
|
||||
|
||||
# This dict keeps track of cached actor tasks. We can't schedule actor
|
||||
# tasks before the actor is actually scheduled/live. So when the caller
|
||||
# tries to schedule a task, we cache it here, and schedule it once the
|
||||
# actor is started.
|
||||
self._pending_actors_to_enqueued_actor_tasks: Dict[
|
||||
TrackedActor, List[Tuple[TrackedActorTask, str, Tuple[Any], Dict[str, Any]]]
|
||||
] = defaultdict(list)
|
||||
|
||||
# ---
|
||||
# Live actors.
|
||||
# We keep one dict for actors that are currently running and a set of
|
||||
# actors that we should forcefully kill.
|
||||
|
||||
# This dict associates the TrackedActor object with the Ray actor handle
|
||||
# and the resources associated to the actor. We use it to schedule the
|
||||
# actual ray tasks, and to return the resources when the actor stopped.
|
||||
self._live_actors_to_ray_actors_resources: Dict[
|
||||
TrackedActor, Tuple[ray.actor.ActorHandle, AcquiredResources]
|
||||
] = {}
|
||||
self._live_resource_cache: Optional[Dict[str, Any]] = None
|
||||
|
||||
# This dict contains all actors that should be killed (after calling
|
||||
# `remove_actor()`). Kill requests will be handled in wait().
|
||||
self._live_actors_to_kill: Set[TrackedActor] = set()
|
||||
|
||||
# Track failed actors
|
||||
self._failed_actor_ids: Set[int] = set()
|
||||
|
||||
def next(self, timeout: Optional[Union[int, float]] = None) -> bool:
|
||||
"""Yield control to event manager to await the next event and invoke callbacks.
|
||||
|
||||
Calling this method will wait for up to ``timeout`` seconds for the next
|
||||
event to arrive.
|
||||
|
||||
When events arrive, callbacks relating to the events will be
|
||||
invoked. A timeout of ``None`` will block until the next event arrives.
|
||||
|
||||
Note:
|
||||
If an actor task fails with a ``RayActorError``, this is one event,
|
||||
but it may trigger _two_ `on_error` callbacks: One for the actor,
|
||||
and one for the task.
|
||||
|
||||
Note:
|
||||
The ``timeout`` argument is used for pure waiting time for events. It does
|
||||
not include time spent on processing callbacks. Depending on the processing
|
||||
time of the callbacks, it can take much longer for this function to
|
||||
return than the specified timeout.
|
||||
|
||||
Args:
|
||||
timeout: Timeout in seconds to wait for next event.
|
||||
|
||||
Returns:
|
||||
True if at least one event was processed.
|
||||
|
||||
"""
|
||||
# First issue any pending forceful actor kills
|
||||
actor_killed = self._try_kill_actor()
|
||||
|
||||
# We always try to start actors as this won't trigger an event callback
|
||||
self._try_start_actors()
|
||||
|
||||
# If an actor was killed, this was our event, and we return.
|
||||
if actor_killed:
|
||||
return True
|
||||
|
||||
# Otherwise, collect all futures and await the next.
|
||||
resource_futures = self._resource_manager.get_resource_futures()
|
||||
actor_state_futures = self._actor_state_events.get_futures()
|
||||
actor_task_futures = self._actor_task_events.get_futures()
|
||||
|
||||
# Shuffle state futures
|
||||
shuffled_state_futures = list(actor_state_futures)
|
||||
random.shuffle(shuffled_state_futures)
|
||||
|
||||
# Shuffle task futures
|
||||
shuffled_task_futures = list(actor_task_futures)
|
||||
random.shuffle(shuffled_task_futures)
|
||||
|
||||
# Prioritize resource futures over actor state over task futures
|
||||
all_futures = resource_futures + shuffled_state_futures + shuffled_task_futures
|
||||
|
||||
start_wait = time.monotonic()
|
||||
ready, _ = ray.wait(all_futures, num_returns=1, timeout=timeout)
|
||||
|
||||
if not ready:
|
||||
return False
|
||||
|
||||
[future] = ready
|
||||
|
||||
if future in actor_state_futures:
|
||||
self._actor_state_events.resolve_future(future)
|
||||
elif future in actor_task_futures:
|
||||
self._actor_task_events.resolve_future(future)
|
||||
else:
|
||||
self._handle_ready_resource_future()
|
||||
# Ready resource futures don't count as one event as they don't trigger
|
||||
# any callbacks. So we repeat until we hit anything that is not a resource
|
||||
# future.
|
||||
time_taken = time.monotonic() - start_wait
|
||||
return self.next(
|
||||
timeout=max(1e-9, timeout - time_taken) if timeout is not None else None
|
||||
)
|
||||
|
||||
self._try_start_actors()
|
||||
return True
|
||||
|
||||
def _actor_start_resolved(self, tracked_actor: TrackedActor, future: ray.ObjectRef):
|
||||
"""Callback to be invoked when actor started"""
|
||||
self._tracked_actors_to_state_futures[tracked_actor].remove(future)
|
||||
|
||||
if tracked_actor._on_start:
|
||||
tracked_actor._on_start(tracked_actor)
|
||||
|
||||
def _actor_stop_resolved(self, tracked_actor: TrackedActor):
|
||||
"""Callback to be invoked when actor stopped"""
|
||||
self._cleanup_actor(tracked_actor=tracked_actor)
|
||||
|
||||
if tracked_actor._on_stop:
|
||||
tracked_actor._on_stop(tracked_actor)
|
||||
|
||||
def _actor_start_failed(self, tracked_actor: TrackedActor, exception: Exception):
|
||||
"""Callback to be invoked when actor start/stop failed"""
|
||||
self._failed_actor_ids.add(tracked_actor.actor_id)
|
||||
|
||||
self._cleanup_actor(tracked_actor=tracked_actor)
|
||||
|
||||
if tracked_actor._on_error:
|
||||
tracked_actor._on_error(tracked_actor, exception)
|
||||
|
||||
def _actor_task_failed(
|
||||
self, tracked_actor_task: TrackedActorTask, exception: Exception
|
||||
):
|
||||
"""Handle an actor task future that became ready.
|
||||
|
||||
- On actor error, trigger actor error callback AND error task error callback
|
||||
- On task error, trigger actor task error callback
|
||||
- On success, trigger actor task result callback
|
||||
"""
|
||||
tracked_actor = tracked_actor_task._tracked_actor
|
||||
|
||||
if isinstance(exception, RayActorError):
|
||||
self._failed_actor_ids.add(tracked_actor.actor_id)
|
||||
|
||||
# Clean up any references to the actor and its futures
|
||||
self._cleanup_actor(tracked_actor=tracked_actor)
|
||||
|
||||
# Handle actor state callbacks
|
||||
if tracked_actor._on_error:
|
||||
tracked_actor._on_error(tracked_actor, exception)
|
||||
|
||||
# Then trigger actor task error callback
|
||||
if tracked_actor_task._on_error:
|
||||
tracked_actor_task._on_error(tracked_actor, exception)
|
||||
|
||||
elif isinstance(exception, RayTaskError):
|
||||
# Otherwise only the task failed. Invoke callback
|
||||
if tracked_actor_task._on_error:
|
||||
tracked_actor_task._on_error(tracked_actor, exception)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Caught unexpected exception: {exception}"
|
||||
) from exception
|
||||
|
||||
def _actor_task_resolved(self, tracked_actor_task: TrackedActorTask, result: Any):
|
||||
tracked_actor = tracked_actor_task._tracked_actor
|
||||
|
||||
# Trigger actor task result callback
|
||||
if tracked_actor_task._on_result:
|
||||
tracked_actor_task._on_result(tracked_actor, result)
|
||||
|
||||
def _handle_ready_resource_future(self):
|
||||
"""Handle a resource future that became ready.
|
||||
|
||||
- Update state of the resource manager
|
||||
- Try to start one actor
|
||||
"""
|
||||
# Force resource manager to update internal state
|
||||
self._resource_manager.update_state()
|
||||
# We handle resource futures one by one, so only try to start 1 actor at a time
|
||||
self._try_start_actors(max_actors=1)
|
||||
|
||||
def _try_start_actors(self, max_actors: Optional[int] = None) -> int:
|
||||
"""Try to start up to ``max_actors`` actors.
|
||||
|
||||
This function will iterate through all resource requests we collected for
|
||||
pending actors. As long as a resource request can be fulfilled (resources
|
||||
are available), we try to start as many actors as possible.
|
||||
|
||||
This will schedule a `Actor.__ray_ready__()` future which, once resolved,
|
||||
will trigger the `TrackedActor.on_start` callback.
|
||||
"""
|
||||
started_actors = 0
|
||||
|
||||
# Iterate through all resource requests
|
||||
for resource_request in self._resource_request_to_pending_actors:
|
||||
if max_actors is not None and started_actors >= max_actors:
|
||||
break
|
||||
|
||||
# While we have resources ready and there are actors left to schedule
|
||||
while (
|
||||
self._resource_manager.has_resources_ready(resource_request)
|
||||
and self._resource_request_to_pending_actors[resource_request]
|
||||
):
|
||||
# Acquire resources for actor
|
||||
acquired_resources = self._resource_manager.acquire_resources(
|
||||
resource_request
|
||||
)
|
||||
assert acquired_resources
|
||||
|
||||
# Get tracked actor to start
|
||||
candidate_actors = self._resource_request_to_pending_actors[
|
||||
resource_request
|
||||
]
|
||||
assert candidate_actors
|
||||
|
||||
tracked_actor = candidate_actors.pop(0)
|
||||
|
||||
# Get actor class and arguments
|
||||
actor_cls, kwargs, _ = self._pending_actors_to_attrs.pop(tracked_actor)
|
||||
|
||||
if not isinstance(actor_cls, ray.actor.ActorClass):
|
||||
actor_cls = ray.remote(actor_cls)
|
||||
|
||||
# Associate to acquired resources
|
||||
[remote_actor_cls] = acquired_resources.annotate_remote_entities(
|
||||
[actor_cls]
|
||||
)
|
||||
|
||||
# Start Ray actor
|
||||
actor = remote_actor_cls.remote(**kwargs)
|
||||
|
||||
# Track
|
||||
self._live_actors_to_ray_actors_resources[tracked_actor] = (
|
||||
actor,
|
||||
acquired_resources,
|
||||
)
|
||||
self._live_resource_cache = None
|
||||
|
||||
# Schedule ready future
|
||||
future = actor.__ray_ready__.remote()
|
||||
|
||||
self._tracked_actors_to_state_futures[tracked_actor].add(future)
|
||||
|
||||
# We need to create the callbacks in a function so tracked_actors
|
||||
# are captured correctly.
|
||||
def create_callbacks(
|
||||
tracked_actor: TrackedActor, future: ray.ObjectRef
|
||||
):
|
||||
def on_actor_start(result: Any):
|
||||
self._actor_start_resolved(
|
||||
tracked_actor=tracked_actor, future=future
|
||||
)
|
||||
|
||||
def on_error(exception: Exception):
|
||||
self._actor_start_failed(
|
||||
tracked_actor=tracked_actor, exception=exception
|
||||
)
|
||||
|
||||
return on_actor_start, on_error
|
||||
|
||||
on_actor_start, on_error = create_callbacks(
|
||||
tracked_actor=tracked_actor, future=future
|
||||
)
|
||||
|
||||
self._actor_state_events.track_future(
|
||||
future=future,
|
||||
on_result=on_actor_start,
|
||||
on_error=on_error,
|
||||
)
|
||||
|
||||
self._enqueue_cached_actor_tasks(tracked_actor=tracked_actor)
|
||||
|
||||
started_actors += 1
|
||||
|
||||
return started_actors
|
||||
|
||||
def _enqueue_cached_actor_tasks(self, tracked_actor: TrackedActor):
|
||||
assert tracked_actor in self._live_actors_to_ray_actors_resources
|
||||
|
||||
# Enqueue cached futures
|
||||
cached_tasks = self._pending_actors_to_enqueued_actor_tasks.pop(
|
||||
tracked_actor, []
|
||||
)
|
||||
for tracked_actor_task, method_name, args, kwargs in cached_tasks:
|
||||
self._schedule_tracked_actor_task(
|
||||
tracked_actor_task=tracked_actor_task,
|
||||
method_name=method_name,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
def _try_kill_actor(self) -> bool:
|
||||
"""Try to kill actor scheduled for termination."""
|
||||
if not self._live_actors_to_kill:
|
||||
return False
|
||||
|
||||
tracked_actor = self._live_actors_to_kill.pop()
|
||||
|
||||
# Remove from tracked actors
|
||||
(
|
||||
ray_actor,
|
||||
acquired_resources,
|
||||
) = self._live_actors_to_ray_actors_resources[tracked_actor]
|
||||
|
||||
# Hard kill if requested
|
||||
ray.kill(ray_actor)
|
||||
|
||||
self._cleanup_actor_futures(tracked_actor)
|
||||
|
||||
self._actor_stop_resolved(tracked_actor)
|
||||
|
||||
return True
|
||||
|
||||
def _cleanup_actor(self, tracked_actor: TrackedActor):
|
||||
self._cleanup_actor_futures(tracked_actor)
|
||||
|
||||
# Remove from tracked actors
|
||||
(
|
||||
ray_actor,
|
||||
acquired_resources,
|
||||
) = self._live_actors_to_ray_actors_resources.pop(tracked_actor)
|
||||
self._live_resource_cache = None
|
||||
|
||||
# Return resources
|
||||
self._resource_manager.free_resources(acquired_resource=acquired_resources)
|
||||
|
||||
@property
|
||||
def all_actors(self) -> List[TrackedActor]:
|
||||
"""Return all ``TrackedActor`` objects managed by this manager instance."""
|
||||
return self.live_actors + self.pending_actors
|
||||
|
||||
@property
|
||||
def live_actors(self) -> List[TrackedActor]:
|
||||
"""Return all ``TrackedActor`` objects that are currently alive."""
|
||||
return list(self._live_actors_to_ray_actors_resources)
|
||||
|
||||
@property
|
||||
def pending_actors(self) -> List[TrackedActor]:
|
||||
"""Return all ``TrackedActor`` objects that are currently pending."""
|
||||
return list(self._pending_actors_to_attrs)
|
||||
|
||||
@property
|
||||
def num_live_actors(self):
|
||||
"""Return number of started actors."""
|
||||
return len(self.live_actors)
|
||||
|
||||
@property
|
||||
def num_pending_actors(self) -> int:
|
||||
"""Return number of pending (not yet started) actors."""
|
||||
return len(self.pending_actors)
|
||||
|
||||
@property
|
||||
def num_total_actors(self):
|
||||
"""Return number of total actors."""
|
||||
return len(self.all_actors)
|
||||
|
||||
@property
|
||||
def num_actor_tasks(self):
|
||||
"""Return number of pending tasks"""
|
||||
return self._actor_task_events.num_futures
|
||||
|
||||
def get_live_actors_resources(self):
|
||||
if self._live_resource_cache:
|
||||
return self._live_resource_cache
|
||||
|
||||
counter = Counter()
|
||||
for _, acq in self._live_actors_to_ray_actors_resources.values():
|
||||
for bdl in acq.resource_request.bundles:
|
||||
counter.update(bdl)
|
||||
self._live_resource_cache = dict(counter)
|
||||
return self._live_resource_cache
|
||||
|
||||
def add_actor(
|
||||
self,
|
||||
cls: Union[Type, ray.actor.ActorClass],
|
||||
kwargs: Dict[str, Any],
|
||||
resource_request: ResourceRequest,
|
||||
*,
|
||||
on_start: Optional[Callable[[TrackedActor], None]] = None,
|
||||
on_stop: Optional[Callable[[TrackedActor], None]] = None,
|
||||
on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
|
||||
) -> TrackedActor:
|
||||
"""Add an actor to be tracked.
|
||||
|
||||
This method will request resources to start the actor. Once the resources
|
||||
are available, the actor will be started and the
|
||||
:meth:`TrackedActor.on_start
|
||||
<ray.air.execution._internal.tracked_actor.TrackedActor.on_start>` callback
|
||||
will be invoked.
|
||||
|
||||
Args:
|
||||
cls: Actor class to schedule.
|
||||
kwargs: Keyword arguments to pass to actor class on construction.
|
||||
resource_request: Resources required to start the actor.
|
||||
on_start: Callback to invoke when the actor started.
|
||||
on_stop: Callback to invoke when the actor stopped.
|
||||
on_error: Callback to invoke when the actor failed.
|
||||
|
||||
Returns:
|
||||
Tracked actor object to reference actor in subsequent API calls.
|
||||
|
||||
"""
|
||||
tracked_actor = TrackedActor(
|
||||
uuid.uuid4().int, on_start=on_start, on_stop=on_stop, on_error=on_error
|
||||
)
|
||||
|
||||
self._pending_actors_to_attrs[tracked_actor] = cls, kwargs, resource_request
|
||||
self._resource_request_to_pending_actors[resource_request].append(tracked_actor)
|
||||
|
||||
self._resource_manager.request_resources(resource_request=resource_request)
|
||||
|
||||
return tracked_actor
|
||||
|
||||
def remove_actor(
|
||||
self,
|
||||
tracked_actor: TrackedActor,
|
||||
kill: bool = False,
|
||||
stop_future: Optional[ray.ObjectRef] = None,
|
||||
) -> bool:
|
||||
"""Remove a tracked actor.
|
||||
|
||||
If the actor has already been started, this will stop the actor. This will
|
||||
trigger the :meth:`TrackedActor.on_stop
|
||||
<ray.air.execution._internal.tracked_actor.TrackedActor.on_stop>`
|
||||
callback once the actor stopped.
|
||||
|
||||
If the actor has only been requested, but not started, yet, this will cancel
|
||||
the actor request. This will not trigger any callback.
|
||||
|
||||
If ``kill=True``, this will use ``ray.kill()`` to forcefully terminate the
|
||||
actor. Otherwise, graceful actor deconstruction will be scheduled after
|
||||
all currently tracked futures are resolved.
|
||||
|
||||
This method returns a boolean, indicating if a stop future is tracked and
|
||||
the ``on_stop`` callback will be invoked. If the actor has been alive,
|
||||
this will be ``True``. If the actor hasn't been scheduled, yet, or failed
|
||||
(and triggered the ``on_error`` callback), this will be ``False``.
|
||||
|
||||
Args:
|
||||
tracked_actor: Tracked actor to be removed.
|
||||
kill: If set, will forcefully terminate the actor instead of gracefully
|
||||
scheduling termination.
|
||||
stop_future: If set, use this future to track actor termination.
|
||||
Otherwise, schedule a ``__ray_terminate__`` future.
|
||||
|
||||
Returns:
|
||||
Boolean indicating if the actor was previously alive, and thus whether
|
||||
a callback will be invoked once it is terminated.
|
||||
|
||||
"""
|
||||
if tracked_actor.actor_id in self._failed_actor_ids:
|
||||
logger.debug(
|
||||
f"Tracked actor already failed, no need to remove: {tracked_actor}"
|
||||
)
|
||||
return False
|
||||
elif tracked_actor in self._live_actors_to_ray_actors_resources:
|
||||
# Ray actor is running.
|
||||
|
||||
if not kill:
|
||||
# Schedule __ray_terminate__ future
|
||||
ray_actor, _ = self._live_actors_to_ray_actors_resources[tracked_actor]
|
||||
|
||||
# Clear state futures here to avoid resolving __ray_ready__ futures
|
||||
for future in list(
|
||||
self._tracked_actors_to_state_futures[tracked_actor]
|
||||
):
|
||||
self._actor_state_events.discard_future(future)
|
||||
self._tracked_actors_to_state_futures[tracked_actor].remove(future)
|
||||
|
||||
# If the __ray_ready__ future hasn't resolved yet, but we already
|
||||
# scheduled the actor via Actor.remote(), we just want to stop
|
||||
# it but not trigger any callbacks. This is in accordance with
|
||||
# the contract defined in the docstring.
|
||||
tracked_actor._on_start = None
|
||||
tracked_actor._on_stop = None
|
||||
tracked_actor._on_error = None
|
||||
|
||||
def on_actor_stop(*args, **kwargs):
|
||||
self._actor_stop_resolved(tracked_actor=tracked_actor)
|
||||
|
||||
if stop_future:
|
||||
# If the stop future was schedule via the actor manager,
|
||||
# discard (track it as state future instead).
|
||||
self._actor_task_events.discard_future(stop_future)
|
||||
else:
|
||||
stop_future = ray_actor.__ray_terminate__.remote()
|
||||
|
||||
self._actor_state_events.track_future(
|
||||
future=stop_future,
|
||||
on_result=on_actor_stop,
|
||||
on_error=on_actor_stop,
|
||||
)
|
||||
|
||||
self._tracked_actors_to_state_futures[tracked_actor].add(stop_future)
|
||||
else:
|
||||
# kill = True
|
||||
self._live_actors_to_kill.add(tracked_actor)
|
||||
|
||||
return True
|
||||
|
||||
elif tracked_actor in self._pending_actors_to_attrs:
|
||||
# Actor is pending, stop
|
||||
_, _, resource_request = self._pending_actors_to_attrs.pop(tracked_actor)
|
||||
self._resource_request_to_pending_actors[resource_request].remove(
|
||||
tracked_actor
|
||||
)
|
||||
self._resource_manager.cancel_resource_request(
|
||||
resource_request=resource_request
|
||||
)
|
||||
return False
|
||||
else:
|
||||
raise ValueError(f"Unknown tracked actor: {tracked_actor}")
|
||||
|
||||
def is_actor_started(self, tracked_actor: TrackedActor) -> bool:
|
||||
"""Returns True if the actor has been started.
|
||||
|
||||
Args:
|
||||
tracked_actor: Tracked actor object.
|
||||
|
||||
Returns:
|
||||
True if the actor has been started, False otherwise.
|
||||
"""
|
||||
return (
|
||||
tracked_actor in self._live_actors_to_ray_actors_resources
|
||||
and tracked_actor.actor_id not in self._failed_actor_ids
|
||||
)
|
||||
|
||||
def is_actor_failed(self, tracked_actor: TrackedActor) -> bool:
|
||||
return tracked_actor.actor_id in self._failed_actor_ids
|
||||
|
||||
def get_actor_resources(
|
||||
self, tracked_actor: TrackedActor
|
||||
) -> Optional[AcquiredResources]:
|
||||
"""Returns the acquired resources of an actor that has been started.
|
||||
|
||||
This will return ``None`` if the actor has not been started, yet.
|
||||
|
||||
Args:
|
||||
tracked_actor: Tracked actor object.
|
||||
|
||||
Returns:
|
||||
The acquired resources of the actor, or ``None`` if the actor has not
|
||||
been started yet.
|
||||
"""
|
||||
if not self.is_actor_started(tracked_actor):
|
||||
return None
|
||||
|
||||
return self._live_actors_to_ray_actors_resources[tracked_actor][1]
|
||||
|
||||
def schedule_actor_task(
|
||||
self,
|
||||
tracked_actor: TrackedActor,
|
||||
method_name: str,
|
||||
args: Optional[Tuple] = None,
|
||||
kwargs: Optional[Dict] = None,
|
||||
on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
|
||||
on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
|
||||
_return_future: bool = False,
|
||||
) -> Optional[ray.ObjectRef]:
|
||||
"""Schedule and track a task on an actor.
|
||||
|
||||
This method will schedule a remote task ``method_name`` on the
|
||||
``tracked_actor``.
|
||||
|
||||
This method accepts two optional callbacks that will be invoked when
|
||||
their respective events are triggered.
|
||||
|
||||
The ``on_result`` callback is triggered when a task resolves successfully.
|
||||
It should accept two arguments: The actor for which the
|
||||
task resolved, and the result received from the remote call.
|
||||
|
||||
The ``on_error`` callback is triggered when a task fails.
|
||||
It should accept two arguments: The actor for which the
|
||||
task threw an error, and the exception.
|
||||
|
||||
Args:
|
||||
tracked_actor: Actor to schedule task on.
|
||||
method_name: Remote method name to invoke on the actor. If this is
|
||||
e.g. ``foo``, then ``actor.foo.remote(*args, **kwargs)`` will be
|
||||
scheduled.
|
||||
args: Arguments to pass to the task.
|
||||
kwargs: Keyword arguments to pass to the task.
|
||||
on_result: Callback to invoke when the task resolves.
|
||||
on_error: Callback to invoke when the task fails.
|
||||
_return_future: If True, return the scheduled task's ``ObjectRef`` for
|
||||
advanced callers. Defaults to False.
|
||||
|
||||
Raises:
|
||||
ValueError: If the ``tracked_actor`` is not managed by this event manager.
|
||||
|
||||
Returns:
|
||||
The scheduled task's ``ObjectRef`` if ``_return_future`` is True,
|
||||
otherwise ``None``.
|
||||
"""
|
||||
args = args or tuple()
|
||||
kwargs = kwargs or {}
|
||||
|
||||
if tracked_actor.actor_id in self._failed_actor_ids:
|
||||
return
|
||||
|
||||
tracked_actor_task = TrackedActorTask(
|
||||
tracked_actor=tracked_actor, on_result=on_result, on_error=on_error
|
||||
)
|
||||
|
||||
if tracked_actor not in self._live_actors_to_ray_actors_resources:
|
||||
# Actor is not started, yet
|
||||
if tracked_actor not in self._pending_actors_to_attrs:
|
||||
raise ValueError(
|
||||
f"Tracked actor is not managed by this event manager: "
|
||||
f"{tracked_actor}"
|
||||
)
|
||||
|
||||
# Cache tasks for future execution
|
||||
self._pending_actors_to_enqueued_actor_tasks[tracked_actor].append(
|
||||
(tracked_actor_task, method_name, args, kwargs)
|
||||
)
|
||||
else:
|
||||
res = self._schedule_tracked_actor_task(
|
||||
tracked_actor_task=tracked_actor_task,
|
||||
method_name=method_name,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
_return_future=_return_future,
|
||||
)
|
||||
if _return_future:
|
||||
return res[1]
|
||||
|
||||
def _schedule_tracked_actor_task(
|
||||
self,
|
||||
tracked_actor_task: TrackedActorTask,
|
||||
method_name: str,
|
||||
*,
|
||||
args: Optional[Tuple] = None,
|
||||
kwargs: Optional[Dict] = None,
|
||||
_return_future: bool = False,
|
||||
) -> Union[TrackedActorTask, Tuple[TrackedActorTask, ray.ObjectRef]]:
|
||||
tracked_actor = tracked_actor_task._tracked_actor
|
||||
ray_actor, _ = self._live_actors_to_ray_actors_resources[tracked_actor]
|
||||
|
||||
try:
|
||||
remote_fn = getattr(ray_actor, method_name)
|
||||
except AttributeError as e:
|
||||
raise AttributeError(
|
||||
f"Remote function `{method_name}()` does not exist for this actor."
|
||||
) from e
|
||||
|
||||
def on_result(result: Any):
|
||||
self._actor_task_resolved(
|
||||
tracked_actor_task=tracked_actor_task, result=result
|
||||
)
|
||||
|
||||
def on_error(exception: Exception):
|
||||
self._actor_task_failed(
|
||||
tracked_actor_task=tracked_actor_task, exception=exception
|
||||
)
|
||||
|
||||
future = remote_fn.remote(*args, **kwargs)
|
||||
|
||||
self._actor_task_events.track_future(
|
||||
future=future, on_result=on_result, on_error=on_error
|
||||
)
|
||||
|
||||
self._tracked_actors_to_task_futures[tracked_actor].add(future)
|
||||
|
||||
if _return_future:
|
||||
return tracked_actor_task, future
|
||||
|
||||
return tracked_actor_task
|
||||
|
||||
def schedule_actor_tasks(
|
||||
self,
|
||||
tracked_actors: List[TrackedActor],
|
||||
method_name: str,
|
||||
*,
|
||||
args: Optional[Union[Tuple, List[Tuple]]] = None,
|
||||
kwargs: Optional[Union[Dict, List[Dict]]] = None,
|
||||
on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
|
||||
on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
|
||||
) -> None:
|
||||
"""Schedule and track tasks on a list of actors.
|
||||
|
||||
This method will schedule a remote task ``method_name`` on all
|
||||
``tracked_actors``.
|
||||
|
||||
``args`` and ``kwargs`` can be a single tuple/dict, in which case the same
|
||||
(keyword) arguments are passed to all actors. If a list is passed instead,
|
||||
they are mapped to the respective actors. In that case, the list of
|
||||
(keyword) arguments must be the same length as the list of actors.
|
||||
|
||||
This method accepts two optional callbacks that will be invoked when
|
||||
their respective events are triggered.
|
||||
|
||||
The ``on_result`` callback is triggered when a task resolves successfully.
|
||||
It should accept two arguments: The actor for which the
|
||||
task resolved, and the result received from the remote call.
|
||||
|
||||
The ``on_error`` callback is triggered when a task fails.
|
||||
It should accept two arguments: The actor for which the
|
||||
task threw an error, and the exception.
|
||||
|
||||
Args:
|
||||
tracked_actors: List of actors to schedule tasks on.
|
||||
method_name: Remote actor method to invoke on the actors. If this is
|
||||
e.g. ``foo``, then ``actor.foo.remote(*args, **kwargs)`` will be
|
||||
scheduled on all actors.
|
||||
args: Arguments to pass to the task.
|
||||
kwargs: Keyword arguments to pass to the task.
|
||||
on_result: Callback to invoke when the task resolves.
|
||||
on_error: Callback to invoke when the task fails.
|
||||
|
||||
"""
|
||||
if not isinstance(args, List):
|
||||
args_list = [args] * len(tracked_actors)
|
||||
else:
|
||||
if len(tracked_actors) != len(args):
|
||||
raise ValueError(
|
||||
f"Length of args must be the same as tracked_actors "
|
||||
f"list. Got `len(kwargs)={len(kwargs)}` and "
|
||||
f"`len(tracked_actors)={len(tracked_actors)}"
|
||||
)
|
||||
args_list = args
|
||||
|
||||
if not isinstance(kwargs, List):
|
||||
kwargs_list = [kwargs] * len(tracked_actors)
|
||||
else:
|
||||
if len(tracked_actors) != len(kwargs):
|
||||
raise ValueError(
|
||||
f"Length of kwargs must be the same as tracked_actors "
|
||||
f"list. Got `len(args)={len(args)}` and "
|
||||
f"`len(tracked_actors)={len(tracked_actors)}"
|
||||
)
|
||||
kwargs_list = kwargs
|
||||
|
||||
for tracked_actor, args, kwargs in zip(tracked_actors, args_list, kwargs_list):
|
||||
self.schedule_actor_task(
|
||||
tracked_actor=tracked_actor,
|
||||
method_name=method_name,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
on_result=on_result,
|
||||
on_error=on_error,
|
||||
)
|
||||
|
||||
def clear_actor_task_futures(self, tracked_actor: TrackedActor):
|
||||
"""Discard all actor task futures from a tracked actor."""
|
||||
futures = self._tracked_actors_to_task_futures.pop(tracked_actor, [])
|
||||
for future in futures:
|
||||
self._actor_task_events.discard_future(future)
|
||||
|
||||
def _cleanup_actor_futures(self, tracked_actor: TrackedActor):
|
||||
# Remove all actor task futures
|
||||
self.clear_actor_task_futures(tracked_actor=tracked_actor)
|
||||
|
||||
# Remove all actor state futures
|
||||
futures = self._tracked_actors_to_state_futures.pop(tracked_actor, [])
|
||||
for future in futures:
|
||||
self._actor_state_events.discard_future(future)
|
||||
|
||||
def cleanup(self):
|
||||
for (
|
||||
actor,
|
||||
acquired_resources,
|
||||
) in self._live_actors_to_ray_actors_resources.values():
|
||||
ray.kill(actor)
|
||||
self._resource_manager.free_resources(acquired_resources)
|
||||
|
||||
for (
|
||||
resource_request,
|
||||
pending_actors,
|
||||
) in self._resource_request_to_pending_actors.items():
|
||||
for i in range(len(pending_actors)):
|
||||
self._resource_manager.cancel_resource_request(resource_request)
|
||||
|
||||
self._resource_manager.clear()
|
||||
|
||||
self.__init__(resource_manager=self._resource_manager)
|
||||
@@ -0,0 +1,93 @@
|
||||
from typing import Any, Callable, List, Optional, Tuple
|
||||
|
||||
|
||||
class Barrier:
|
||||
"""Barrier to collect results and process them in bulk.
|
||||
|
||||
A barrier can be used to collect multiple results and process them in bulk once
|
||||
a certain count or a timeout is reached.
|
||||
|
||||
For instance, if ``max_results=N``, the ``on_completion`` callback will be
|
||||
invoked once :meth:`arrive` has been called ``N`` times.
|
||||
|
||||
The completion callback will only be invoked once, even if more results
|
||||
arrive after completion. The collected results can be resetted
|
||||
with :meth:`reset`, after which the callback may be invoked again.
|
||||
|
||||
The completion callback should expect one argument, which is the barrier
|
||||
object that completed.
|
||||
|
||||
Args:
|
||||
max_results: Maximum number of results to collect before a call to
|
||||
:meth:`wait` resolves or the :meth:`on_completion` callback is invoked.
|
||||
on_completion: Callback to invoke when ``max_results`` results
|
||||
arrived at the barrier.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_results: int,
|
||||
*,
|
||||
on_completion: Optional[Callable[["Barrier"], None]] = None,
|
||||
):
|
||||
self._max_results = max_results
|
||||
|
||||
# on_completion callback
|
||||
self._completed = False
|
||||
self._on_completion = on_completion
|
||||
|
||||
# Collect received results
|
||||
self._results: List[Tuple[Any]] = []
|
||||
|
||||
def arrive(self, *data: Any):
|
||||
"""Notify barrier that a result successfully arrived.
|
||||
|
||||
This will count against the ``max_results`` limit. The received result
|
||||
will be included in a call to :meth:`get_results`.
|
||||
|
||||
Args:
|
||||
*data: Result data to be cached. Can be obtained via :meth:`get_results`.
|
||||
|
||||
"""
|
||||
if len(data) == 1:
|
||||
data = data[0]
|
||||
|
||||
self._results.append(data)
|
||||
self._check_completion()
|
||||
|
||||
def _check_completion(self):
|
||||
if self._completed:
|
||||
# Already fired completion callback
|
||||
return
|
||||
|
||||
if self.num_results >= self._max_results:
|
||||
# Barrier is complete
|
||||
self._completed = True
|
||||
|
||||
if self._on_completion:
|
||||
self._on_completion(self)
|
||||
|
||||
@property
|
||||
def completed(self) -> bool:
|
||||
"""Returns True if the barrier is completed."""
|
||||
return self._completed
|
||||
|
||||
@property
|
||||
def num_results(self) -> int:
|
||||
"""Number of received (successful) results."""
|
||||
return len(self._results)
|
||||
|
||||
def get_results(self) -> List[Tuple[Any]]:
|
||||
"""Return list of received results."""
|
||||
return self._results
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset barrier, removing all received results.
|
||||
|
||||
Resetting the barrier will reset the completion status. When ``max_results``
|
||||
is set and enough new events arrive after resetting, the
|
||||
:meth:`on_completion` callback will be invoked again.
|
||||
"""
|
||||
self._completed = False
|
||||
self._results = []
|
||||
@@ -0,0 +1,148 @@
|
||||
import random
|
||||
from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Union
|
||||
|
||||
import ray
|
||||
|
||||
_ResultCallback = Callable[[Any], None]
|
||||
_ErrorCallback = Callable[[Exception], None]
|
||||
|
||||
|
||||
class RayEventManager:
|
||||
"""Event manager for Ray futures.
|
||||
|
||||
The event manager can be used to track futures and invoke callbacks when
|
||||
they resolve.
|
||||
|
||||
Futures are tracked with :meth:`track_future`. Future can then be awaited with
|
||||
:meth:`wait`. When futures successfully resolve, they trigger an optional
|
||||
``on_result`` callback that can be passed to :meth:`track_future`. If they
|
||||
fail, they trigger an optional ``on_error`` callback.
|
||||
|
||||
Args:
|
||||
shuffle_futures: If True, futures will be shuffled before awaited. This
|
||||
will avoid implicit prioritization of futures within Ray.
|
||||
"""
|
||||
|
||||
def __init__(self, shuffle_futures: bool = True):
|
||||
self._shuffle_futures = shuffle_futures
|
||||
|
||||
# Map of futures to callbacks (result, error)
|
||||
self._tracked_futures: Dict[
|
||||
ray.ObjectRef, Tuple[Optional[_ResultCallback], Optional[_ErrorCallback]]
|
||||
] = {}
|
||||
|
||||
def track_future(
|
||||
self,
|
||||
future: ray.ObjectRef,
|
||||
on_result: Optional[_ResultCallback] = None,
|
||||
on_error: Optional[_ErrorCallback] = None,
|
||||
):
|
||||
"""Track a single future and invoke callbacks on resolution.
|
||||
|
||||
Control has to be yielded to the event manager for the callbacks to
|
||||
be invoked, either via :meth:`wait` or via :meth:`resolve_future`.
|
||||
|
||||
Args:
|
||||
future: Ray future to await.
|
||||
on_result: Callback to invoke when the future resolves successfully.
|
||||
on_error: Callback to invoke when the future fails.
|
||||
|
||||
"""
|
||||
self._tracked_futures[future] = (on_result, on_error)
|
||||
|
||||
def track_futures(
|
||||
self,
|
||||
futures: Iterable[ray.ObjectRef],
|
||||
on_result: Optional[_ResultCallback] = None,
|
||||
on_error: Optional[_ErrorCallback] = None,
|
||||
):
|
||||
"""Track multiple futures and invoke callbacks on resolution.
|
||||
|
||||
Control has to be yielded to the event manager for the callbacks to
|
||||
be invoked, either via :meth:`wait` or via :meth:`resolve_future`.
|
||||
|
||||
Args:
|
||||
futures: Ray futures to await.
|
||||
on_result: Callback to invoke when the future resolves successfully.
|
||||
on_error: Callback to invoke when the future fails.
|
||||
|
||||
"""
|
||||
for future in futures:
|
||||
self.track_future(future, on_result=on_result, on_error=on_error)
|
||||
|
||||
def discard_future(self, future: ray.ObjectRef):
|
||||
"""Remove future from tracking.
|
||||
|
||||
The future will not be awaited anymore, and it will not trigger any callbacks.
|
||||
|
||||
Args:
|
||||
future: Ray futures to discard.
|
||||
"""
|
||||
self._tracked_futures.pop(future, None)
|
||||
|
||||
def get_futures(self) -> Set[ray.ObjectRef]:
|
||||
"""Get futures tracked by the event manager."""
|
||||
return set(self._tracked_futures)
|
||||
|
||||
@property
|
||||
def num_futures(self) -> int:
|
||||
return len(self._tracked_futures)
|
||||
|
||||
def resolve_future(self, future: ray.ObjectRef):
|
||||
"""Resolve a single future.
|
||||
|
||||
This method will block until the future is available. It will then
|
||||
trigger the callback associated to the future and the event (success
|
||||
or error), if specified.
|
||||
|
||||
Args:
|
||||
future: Ray future to resolve.
|
||||
|
||||
"""
|
||||
try:
|
||||
on_result, on_error = self._tracked_futures.pop(future)
|
||||
except KeyError as e:
|
||||
raise ValueError(
|
||||
f"Future {future} is not tracked by this RayEventManager"
|
||||
) from e
|
||||
|
||||
try:
|
||||
result = ray.get(future)
|
||||
except Exception as e:
|
||||
if on_error:
|
||||
on_error(e)
|
||||
else:
|
||||
raise e
|
||||
else:
|
||||
if on_result:
|
||||
on_result(result)
|
||||
|
||||
def wait(
|
||||
self,
|
||||
timeout: Optional[Union[float, int]] = None,
|
||||
num_results: Optional[int] = 1,
|
||||
):
|
||||
"""Wait up to ``timeout`` seconds for ``num_results`` futures to resolve.
|
||||
|
||||
If ``timeout=None``, this method will block until all `num_results`` futures
|
||||
resolve. If ``num_results=None``, this method will await all tracked futures.
|
||||
|
||||
For every future that resolves, the respective associated callbacks will be
|
||||
invoked.
|
||||
|
||||
Args:
|
||||
timeout: Timeout in second to wait for futures to resolve.
|
||||
num_results: Number of futures to await. If ``None``, will wait for
|
||||
all tracked futures to resolve.
|
||||
|
||||
"""
|
||||
futures = list(self.get_futures())
|
||||
|
||||
if self._shuffle_futures:
|
||||
random.shuffle(futures)
|
||||
|
||||
num_results = num_results or len(futures)
|
||||
|
||||
ready, _ = ray.wait(list(futures), timeout=timeout, num_returns=num_results)
|
||||
for future in ready:
|
||||
self.resolve_future(future)
|
||||
@@ -0,0 +1,62 @@
|
||||
from typing import Callable, Optional
|
||||
|
||||
|
||||
class TrackedActor:
|
||||
"""Actor tracked by an actor manager.
|
||||
|
||||
This object is used to reference a Ray actor on an actor manager
|
||||
|
||||
Existence of this object does not mean that the Ray actor has already been started.
|
||||
Actor state can be inquired from the actor manager tracking the Ray actor.
|
||||
|
||||
Note:
|
||||
Objects of this class are returned by the :class:`RayActorManager`.
|
||||
This class should not be instantiated manually.
|
||||
|
||||
Attributes:
|
||||
actor_id: ID for identification of the actor within the actor manager. This
|
||||
ID is not related to the Ray actor ID.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
actor_id: int,
|
||||
on_start: Optional[Callable[["TrackedActor"], None]] = None,
|
||||
on_stop: Optional[Callable[["TrackedActor"], None]] = None,
|
||||
on_error: Optional[Callable[["TrackedActor", Exception], None]] = None,
|
||||
):
|
||||
"""Initialize the tracked actor.
|
||||
|
||||
Args:
|
||||
actor_id: ID for identification of the actor within the actor manager.
|
||||
on_start: Callback to invoke when the actor started.
|
||||
on_stop: Callback to invoke when the actor stopped.
|
||||
on_error: Callback to invoke when the actor failed.
|
||||
"""
|
||||
self.actor_id = actor_id
|
||||
self._on_start = on_start
|
||||
self._on_stop = on_stop
|
||||
self._on_error = on_error
|
||||
|
||||
def set_on_start(self, on_start: Optional[Callable[["TrackedActor"], None]]):
|
||||
self._on_start = on_start
|
||||
|
||||
def set_on_stop(self, on_stop: Optional[Callable[["TrackedActor"], None]]):
|
||||
self._on_stop = on_stop
|
||||
|
||||
def set_on_error(
|
||||
self, on_error: Optional[Callable[["TrackedActor", Exception], None]]
|
||||
):
|
||||
self._on_error = on_error
|
||||
|
||||
def __repr__(self):
|
||||
return f"<TrackedActor {self.actor_id}>"
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, self.__class__):
|
||||
return False
|
||||
return self.actor_id == other.actor_id
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.actor_id)
|
||||
@@ -0,0 +1,42 @@
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from ray.air.execution._internal.tracked_actor import TrackedActor
|
||||
|
||||
|
||||
class TrackedActorTask:
|
||||
"""Actor task tracked by a Ray event manager.
|
||||
|
||||
This container class is used to define callbacks to be invoked when
|
||||
the task resolves, errors, or times out.
|
||||
|
||||
Note:
|
||||
Objects of this class are returned by the :class:`RayActorManager`.
|
||||
This class should not be instantiated manually.
|
||||
|
||||
Args:
|
||||
tracked_actor: Tracked actor object this task is scheduled on.
|
||||
on_result: Callback to invoke when the task resolves.
|
||||
on_error: Callback to invoke when the task fails.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
tracked_futures = actor_manager.schedule_actor_tasks(
|
||||
actor_manager.live_actors,
|
||||
"foo",
|
||||
on_result=lambda actor, result: print(result)
|
||||
)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tracked_actor: TrackedActor,
|
||||
on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
|
||||
on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
|
||||
):
|
||||
self._tracked_actor = tracked_actor
|
||||
|
||||
self._on_result = on_result
|
||||
self._on_error = on_error
|
||||
@@ -0,0 +1,12 @@
|
||||
from ray.air.execution.resources.fixed import FixedResourceManager
|
||||
from ray.air.execution.resources.placement_group import PlacementGroupResourceManager
|
||||
from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
|
||||
from ray.air.execution.resources.resource_manager import ResourceManager
|
||||
|
||||
__all__ = [
|
||||
"ResourceRequest",
|
||||
"AcquiredResources",
|
||||
"ResourceManager",
|
||||
"FixedResourceManager",
|
||||
"PlacementGroupResourceManager",
|
||||
]
|
||||
@@ -0,0 +1,147 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray import SCRIPT_MODE
|
||||
from ray.air.execution.resources.request import (
|
||||
AcquiredResources,
|
||||
RemoteRayEntity,
|
||||
ResourceRequest,
|
||||
)
|
||||
from ray.air.execution.resources.resource_manager import ResourceManager
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
# Avoid numerical errors by multiplying and subtracting with this number.
|
||||
# Compare: 0.99 - 0.33 = 0.65999... vs (0.99 * 1000 - 0.33 * 1000) / 1000 = 0.66
|
||||
_DIGITS = 100000
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
@dataclass
|
||||
class FixedAcquiredResources(AcquiredResources):
|
||||
bundles: List[Dict[str, float]]
|
||||
|
||||
def _annotate_remote_entity(
|
||||
self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
|
||||
) -> RemoteRayEntity:
|
||||
bundle = bundle.copy()
|
||||
num_cpus = bundle.pop("CPU", 0)
|
||||
num_gpus = bundle.pop("GPU", 0)
|
||||
memory = bundle.pop("memory", 0.0)
|
||||
|
||||
return entity.options(
|
||||
num_cpus=num_cpus,
|
||||
num_gpus=num_gpus,
|
||||
memory=memory,
|
||||
resources=bundle,
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class FixedResourceManager(ResourceManager):
|
||||
"""Fixed budget based resource manager.
|
||||
|
||||
This resource manager keeps track of a fixed set of resources. When resources
|
||||
are acquired, they are subtracted from the budget. When resources are freed,
|
||||
they are added back to the budget.
|
||||
|
||||
The resource manager still requires resources to be requested before they become
|
||||
available. However, because the resource requests are virtual, this will not
|
||||
trigger autoscaling.
|
||||
|
||||
Additionally, resources are not reserved on request, only on acquisition. Thus,
|
||||
acquiring a resource can change the availability of other requests. Note that
|
||||
this behavior may be changed in future implementations.
|
||||
|
||||
The fixed resource manager does not support placement strategies. Using
|
||||
``STRICT_SPREAD`` will result in an error. ``STRICT_PACK`` will succeed only
|
||||
within a placement group bundle. All other placement group arguments will be
|
||||
ignored.
|
||||
|
||||
Args:
|
||||
total_resources: Budget of resources to manage. Defaults to all available
|
||||
resources in the current task or all cluster resources (if outside a task).
|
||||
|
||||
"""
|
||||
|
||||
_resource_cls: AcquiredResources = FixedAcquiredResources
|
||||
|
||||
def __init__(self, total_resources: Optional[Dict[str, float]] = None):
|
||||
rtc = ray.get_runtime_context()
|
||||
|
||||
if not total_resources:
|
||||
if rtc.worker.mode in {None, SCRIPT_MODE}:
|
||||
total_resources = ray.cluster_resources()
|
||||
else:
|
||||
total_resources = rtc.get_assigned_resources()
|
||||
|
||||
# If we are in a placement group, all of our resources will be in a bundle
|
||||
# and thus fulfill requirements of STRICT_PACK - but only if child tasks
|
||||
# are captured by the pg.
|
||||
self._allow_strict_pack = (
|
||||
ray.util.get_current_placement_group() is not None
|
||||
and rtc.should_capture_child_tasks_in_placement_group
|
||||
)
|
||||
|
||||
self._total_resources = total_resources
|
||||
self._requested_resources = []
|
||||
self._used_resources = []
|
||||
|
||||
@property
|
||||
def _available_resources(self) -> Dict[str, float]:
|
||||
available_resources = self._total_resources.copy()
|
||||
|
||||
for used_resources in self._used_resources:
|
||||
all_resources = used_resources.required_resources
|
||||
for k, v in all_resources.items():
|
||||
available_resources[k] = (
|
||||
available_resources[k] * _DIGITS - v * _DIGITS
|
||||
) / _DIGITS
|
||||
return available_resources
|
||||
|
||||
def request_resources(self, resource_request: ResourceRequest):
|
||||
if resource_request.strategy == "STRICT_SPREAD" or (
|
||||
not self._allow_strict_pack and resource_request.strategy == "STRICT_PACK"
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Requested a resource with placement strategy "
|
||||
f"{resource_request.strategy}, but this cannot be fulfilled by a "
|
||||
f"FixedResourceManager. In a nested setting, please set the inner "
|
||||
f"placement strategy to be less restrictive (i.e. no STRICT_ strategy)."
|
||||
)
|
||||
|
||||
self._requested_resources.append(resource_request)
|
||||
|
||||
def cancel_resource_request(self, resource_request: ResourceRequest):
|
||||
self._requested_resources.remove(resource_request)
|
||||
|
||||
def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
|
||||
if resource_request not in self._requested_resources:
|
||||
return False
|
||||
|
||||
available_resources = self._available_resources
|
||||
all_resources = resource_request.required_resources
|
||||
for k, v in all_resources.items():
|
||||
if available_resources.get(k, 0.0) < v:
|
||||
return False
|
||||
return True
|
||||
|
||||
def acquire_resources(
|
||||
self, resource_request: ResourceRequest
|
||||
) -> Optional[AcquiredResources]:
|
||||
if not self.has_resources_ready(resource_request):
|
||||
return None
|
||||
|
||||
self._used_resources.append(resource_request)
|
||||
return self._resource_cls(
|
||||
bundles=resource_request.bundles, resource_request=resource_request
|
||||
)
|
||||
|
||||
def free_resources(self, acquired_resource: AcquiredResources):
|
||||
resources = acquired_resource.resource_request
|
||||
self._used_resources.remove(resources)
|
||||
|
||||
def clear(self):
|
||||
# Reset internal state
|
||||
self._requested_resources = []
|
||||
self._used_resources = []
|
||||
@@ -0,0 +1,214 @@
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Set
|
||||
|
||||
import ray
|
||||
from ray.air.execution.resources.request import (
|
||||
AcquiredResources,
|
||||
RemoteRayEntity,
|
||||
ResourceRequest,
|
||||
)
|
||||
from ray.air.execution.resources.resource_manager import ResourceManager
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
from ray.util.placement_group import PlacementGroup, remove_placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
@dataclass
|
||||
class PlacementGroupAcquiredResources(AcquiredResources):
|
||||
placement_group: PlacementGroup
|
||||
|
||||
def _annotate_remote_entity(
|
||||
self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
|
||||
) -> RemoteRayEntity:
|
||||
bundle = bundle.copy()
|
||||
num_cpus = bundle.pop("CPU", 0)
|
||||
num_gpus = bundle.pop("GPU", 0)
|
||||
memory = bundle.pop("memory", 0.0)
|
||||
|
||||
return entity.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=self.placement_group,
|
||||
placement_group_bundle_index=bundle_index,
|
||||
placement_group_capture_child_tasks=True,
|
||||
),
|
||||
num_cpus=num_cpus,
|
||||
num_gpus=num_gpus,
|
||||
memory=memory,
|
||||
resources=bundle,
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class PlacementGroupResourceManager(ResourceManager):
|
||||
"""Resource manager using placement groups as the resource backend.
|
||||
|
||||
This manager will use placement groups to fulfill resource requests. Requesting
|
||||
a resource will schedule the placement group. Acquiring a resource will
|
||||
return a ``PlacementGroupAcquiredResources`` that can be used to schedule
|
||||
Ray tasks and actors on the placement group. Freeing an acquired resource
|
||||
will destroy the associated placement group.
|
||||
|
||||
Ray core does not emit events when resources are available. Instead, the
|
||||
scheduling state has to be periodically updated.
|
||||
|
||||
Per default, placement group scheduling state is refreshed every time when
|
||||
resource state is inquired, but not more often than once every ``update_interval_s``
|
||||
seconds. Alternatively, staging futures can be retrieved (and awaited) with
|
||||
``get_resource_futures()`` and state update can be force with ``update_state()``.
|
||||
|
||||
Args:
|
||||
update_interval_s: Minimum interval in seconds between updating scheduling
|
||||
state of placement groups.
|
||||
|
||||
"""
|
||||
|
||||
_resource_cls: AcquiredResources = PlacementGroupAcquiredResources
|
||||
|
||||
def __init__(self, update_interval_s: float = 0.1):
|
||||
# Internally, the placement group lifecycle is like this:
|
||||
# - Resources are requested with ``request_resources()``
|
||||
# - A placement group is scheduled ("staged")
|
||||
# - A ``PlacementGroup.ready()`` future is scheduled ("staging future")
|
||||
# - We update the scheduling state when we need to
|
||||
# (e.g. when ``has_resources_ready()`` is called)
|
||||
# - When staging futures resolve, a placement group is moved from "staging"
|
||||
# to "ready"
|
||||
# - When a resource request is canceled, we remove a placement group from
|
||||
# "staging". If there are not staged placement groups
|
||||
# (because they are already "ready"), we remove one from "ready" instead.
|
||||
# - When a resource is acquired, the pg is removed from "ready" and moved
|
||||
# to "acquired"
|
||||
# - When a resource is freed, the pg is removed from "acquired" and destroyed
|
||||
|
||||
# Mapping of placement group to request
|
||||
self._pg_to_request: Dict[PlacementGroup, ResourceRequest] = {}
|
||||
|
||||
# PGs that are staged but not "ready", yet (i.e. not CREATED)
|
||||
self._request_to_staged_pgs: Dict[
|
||||
ResourceRequest, Set[PlacementGroup]
|
||||
] = defaultdict(set)
|
||||
|
||||
# PGs that are CREATED and can be used by tasks and actors
|
||||
self._request_to_ready_pgs: Dict[
|
||||
ResourceRequest, Set[PlacementGroup]
|
||||
] = defaultdict(set)
|
||||
|
||||
# Staging futures used to update internal state.
|
||||
# We keep a double mapping here for better lookup efficiency.
|
||||
self._staging_future_to_pg: Dict[ray.ObjectRef, PlacementGroup] = dict()
|
||||
self._pg_to_staging_future: Dict[PlacementGroup, ray.ObjectRef] = dict()
|
||||
|
||||
# Set of acquired PGs. We keep track of these here to make sure we
|
||||
# only free PGs that this manager managed.
|
||||
self._acquired_pgs: Set[PlacementGroup] = set()
|
||||
|
||||
# Minimum time between updates of the internal state
|
||||
self.update_interval_s = update_interval_s
|
||||
self._last_update = time.monotonic() - self.update_interval_s - 1
|
||||
|
||||
def get_resource_futures(self) -> List[ray.ObjectRef]:
|
||||
return list(self._staging_future_to_pg.keys())
|
||||
|
||||
def _maybe_update_state(self):
|
||||
now = time.monotonic()
|
||||
if now > self._last_update + self.update_interval_s:
|
||||
self.update_state()
|
||||
|
||||
def update_state(self):
|
||||
ready, not_ready = ray.wait(
|
||||
list(self._staging_future_to_pg.keys()),
|
||||
num_returns=len(self._staging_future_to_pg),
|
||||
timeout=0,
|
||||
)
|
||||
for future in ready:
|
||||
# Remove staging future
|
||||
pg = self._staging_future_to_pg.pop(future)
|
||||
self._pg_to_staging_future.pop(pg)
|
||||
# Fetch resource request
|
||||
request = self._pg_to_request[pg]
|
||||
# Remove from staging, add to ready
|
||||
self._request_to_staged_pgs[request].remove(pg)
|
||||
self._request_to_ready_pgs[request].add(pg)
|
||||
self._last_update = time.monotonic()
|
||||
|
||||
def request_resources(self, resource_request: ResourceRequest):
|
||||
pg = resource_request.to_placement_group()
|
||||
self._pg_to_request[pg] = resource_request
|
||||
self._request_to_staged_pgs[resource_request].add(pg)
|
||||
|
||||
future = pg.ready()
|
||||
self._staging_future_to_pg[future] = pg
|
||||
self._pg_to_staging_future[pg] = future
|
||||
|
||||
def cancel_resource_request(self, resource_request: ResourceRequest):
|
||||
if self._request_to_staged_pgs[resource_request]:
|
||||
pg = self._request_to_staged_pgs[resource_request].pop()
|
||||
|
||||
# PG was staging
|
||||
future = self._pg_to_staging_future.pop(pg)
|
||||
self._staging_future_to_pg.pop(future)
|
||||
|
||||
# Cancel the pg.ready task.
|
||||
# Otherwise, it will be pending node assignment forever.
|
||||
ray.cancel(future)
|
||||
else:
|
||||
# PG might be ready
|
||||
pg = self._request_to_ready_pgs[resource_request].pop()
|
||||
if not pg:
|
||||
raise RuntimeError(
|
||||
"Cannot cancel resource request: No placement group was "
|
||||
f"staged or is ready. Make sure to not cancel more resource "
|
||||
f"requests than you've created. Request: {resource_request}"
|
||||
)
|
||||
|
||||
self._pg_to_request.pop(pg)
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
|
||||
if not bool(len(self._request_to_ready_pgs[resource_request])):
|
||||
# Only update state if needed
|
||||
self._maybe_update_state()
|
||||
|
||||
return bool(len(self._request_to_ready_pgs[resource_request]))
|
||||
|
||||
def acquire_resources(
|
||||
self, resource_request: ResourceRequest
|
||||
) -> Optional[PlacementGroupAcquiredResources]:
|
||||
if not self.has_resources_ready(resource_request):
|
||||
return None
|
||||
|
||||
pg = self._request_to_ready_pgs[resource_request].pop()
|
||||
self._acquired_pgs.add(pg)
|
||||
|
||||
return self._resource_cls(placement_group=pg, resource_request=resource_request)
|
||||
|
||||
def free_resources(self, acquired_resource: PlacementGroupAcquiredResources):
|
||||
pg = acquired_resource.placement_group
|
||||
|
||||
self._acquired_pgs.remove(pg)
|
||||
remove_placement_group(pg)
|
||||
self._pg_to_request.pop(pg)
|
||||
|
||||
def clear(self):
|
||||
if not ray.is_initialized():
|
||||
return
|
||||
|
||||
for staged_pgs in self._request_to_staged_pgs.values():
|
||||
for staged_pg in staged_pgs:
|
||||
remove_placement_group(staged_pg)
|
||||
|
||||
for ready_pgs in self._request_to_ready_pgs.values():
|
||||
for ready_pg in ready_pgs:
|
||||
remove_placement_group(ready_pg)
|
||||
|
||||
for acquired_pg in self._acquired_pgs:
|
||||
remove_placement_group(acquired_pg)
|
||||
|
||||
# Reset internal state
|
||||
self.__init__(update_interval_s=self.update_interval_s)
|
||||
|
||||
def __del__(self):
|
||||
self.clear()
|
||||
@@ -0,0 +1,259 @@
|
||||
import abc
|
||||
import json
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from inspect import signature
|
||||
from typing import Dict, List, Union
|
||||
|
||||
import ray
|
||||
from ray.util import placement_group
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
RemoteRayEntity = Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]
|
||||
|
||||
|
||||
def _sum_bundles(bundles: List[Dict[str, float]]) -> Dict[str, float]:
|
||||
"""Sum all resources in a list of resource bundles.
|
||||
|
||||
Args:
|
||||
bundles: List of resource bundles.
|
||||
|
||||
Returns:
|
||||
Dict containing all resources summed up.
|
||||
"""
|
||||
resources = {}
|
||||
for bundle in bundles:
|
||||
for k, v in bundle.items():
|
||||
resources[k] = resources.get(k, 0) + v
|
||||
return resources
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ResourceRequest:
|
||||
"""Request for resources.
|
||||
|
||||
This class is used to define a resource request. A resource request comprises one
|
||||
or more bundles of resources and instructions on the scheduling behavior.
|
||||
|
||||
The resource request can be submitted to a resource manager, which will
|
||||
schedule the resources. Depending on the resource backend, this may instruct
|
||||
Ray to scale up (autoscaling).
|
||||
|
||||
Resource requests are compatible with the most fine-grained low-level resource
|
||||
backend, which are Ray placement groups.
|
||||
|
||||
Args:
|
||||
bundles: A list of bundles which represent the resources requirements.
|
||||
E.g. ``[{"CPU": 1, "GPU": 1}]``.
|
||||
strategy: The scheduling strategy to acquire the bundles.
|
||||
|
||||
- "PACK": Packs Bundles into as few nodes as possible.
|
||||
- "SPREAD": Places Bundles across distinct nodes as even as possible.
|
||||
- "STRICT_PACK": Packs Bundles into one node. The group is
|
||||
not allowed to span multiple nodes.
|
||||
- "STRICT_SPREAD": Packs Bundles across distinct nodes.
|
||||
*args: Passed to the call of ``placement_group()``, if applicable.
|
||||
**kwargs: Passed to the call of ``placement_group()``, if applicable.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bundles: List[Dict[str, Union[int, float]]],
|
||||
strategy: str = "PACK",
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if not bundles:
|
||||
raise ValueError("Cannot initialize a ResourceRequest with zero bundles.")
|
||||
|
||||
# Remove empty resource keys
|
||||
self._bundles = [
|
||||
{k: float(v) for k, v in bundle.items() if v != 0} for bundle in bundles
|
||||
]
|
||||
|
||||
# Check if the head bundle is empty (no resources defined or all resources
|
||||
# are 0 (and thus removed in the previous step)
|
||||
if not self._bundles[0]:
|
||||
# This is when the head bundle doesn't need resources.
|
||||
self._head_bundle_is_empty = True
|
||||
self._bundles.pop(0)
|
||||
|
||||
if not self._bundles:
|
||||
raise ValueError(
|
||||
"Cannot initialize a ResourceRequest with an empty head "
|
||||
"and zero worker bundles."
|
||||
)
|
||||
else:
|
||||
self._head_bundle_is_empty = False
|
||||
|
||||
self._strategy = strategy
|
||||
self._args = args
|
||||
self._kwargs = kwargs
|
||||
|
||||
self._hash = None
|
||||
self._bound = None
|
||||
|
||||
self._bind()
|
||||
|
||||
@property
|
||||
def head_bundle_is_empty(self):
|
||||
"""Returns True if head bundle is empty while child bundles
|
||||
need resources.
|
||||
|
||||
This is considered an internal API within Tune.
|
||||
"""
|
||||
return self._head_bundle_is_empty
|
||||
|
||||
@property
|
||||
@DeveloperAPI
|
||||
def head_cpus(self) -> float:
|
||||
"""Returns the number of cpus in the head bundle."""
|
||||
return 0.0 if self._head_bundle_is_empty else self._bundles[0].get("CPU", 0.0)
|
||||
|
||||
@property
|
||||
@DeveloperAPI
|
||||
def bundles(self) -> List[Dict[str, float]]:
|
||||
"""Returns a deep copy of resource bundles"""
|
||||
return deepcopy(self._bundles)
|
||||
|
||||
@property
|
||||
def required_resources(self) -> Dict[str, float]:
|
||||
"""Returns a dict containing the sums of all resources"""
|
||||
return _sum_bundles(self._bundles)
|
||||
|
||||
@property
|
||||
@DeveloperAPI
|
||||
def strategy(self) -> str:
|
||||
"""Returns the placement strategy"""
|
||||
return self._strategy
|
||||
|
||||
def _bind(self):
|
||||
"""Bind the args and kwargs to the `placement_group()` signature.
|
||||
|
||||
We bind the args and kwargs, so we can compare equality of two resource
|
||||
requests. The main reason for this is that the `placement_group()` API
|
||||
can evolve independently from the ResourceRequest API (e.g. adding new
|
||||
arguments). Then, `ResourceRequest(bundles, strategy, arg=arg)` should
|
||||
be the same as `ResourceRequest(bundles, strategy, arg)`.
|
||||
"""
|
||||
sig = signature(placement_group)
|
||||
try:
|
||||
self._bound = sig.bind(
|
||||
self._bundles, self._strategy, *self._args, **self._kwargs
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"Invalid definition for resource request. Please check "
|
||||
"that you passed valid arguments to the ResourceRequest "
|
||||
"object."
|
||||
) from exc
|
||||
|
||||
def to_placement_group(self):
|
||||
return placement_group(*self._bound.args, **self._bound.kwargs)
|
||||
|
||||
def __eq__(self, other: "ResourceRequest"):
|
||||
return (
|
||||
isinstance(other, ResourceRequest)
|
||||
and self._bound == other._bound
|
||||
and self.head_bundle_is_empty == other.head_bundle_is_empty
|
||||
)
|
||||
|
||||
def __hash__(self):
|
||||
if not self._hash:
|
||||
# Cache hash
|
||||
self._hash = hash(
|
||||
json.dumps(
|
||||
{"args": self._bound.args, "kwargs": self._bound.kwargs},
|
||||
sort_keys=True,
|
||||
indent=0,
|
||||
ensure_ascii=True,
|
||||
)
|
||||
)
|
||||
return self._hash
|
||||
|
||||
def __getstate__(self):
|
||||
state = self.__dict__.copy()
|
||||
state.pop("_hash", None)
|
||||
state.pop("_bound", None)
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.__dict__.update(state)
|
||||
self._hash = None
|
||||
self._bound = None
|
||||
self._bind()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"<ResourceRequest (_bound={self._bound}, "
|
||||
f"head_bundle_is_empty={self.head_bundle_is_empty})>"
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
@dataclass
|
||||
class AcquiredResources(abc.ABC):
|
||||
"""Base class for resources that have been acquired.
|
||||
|
||||
Acquired resources can be associated to Ray objects, which can then be
|
||||
scheduled using these resources.
|
||||
|
||||
Internally this can point e.g. to a placement group, a placement
|
||||
group bundle index, or just raw resources.
|
||||
|
||||
The main API is the `annotate_remote_entities` method. This will associate
|
||||
remote Ray objects (tasks and actors) with the acquired resources by setting
|
||||
the Ray remote options to use the acquired resources.
|
||||
"""
|
||||
|
||||
resource_request: ResourceRequest
|
||||
|
||||
def annotate_remote_entities(
|
||||
self, entities: List[RemoteRayEntity]
|
||||
) -> List[Union[RemoteRayEntity]]:
|
||||
"""Return remote ray entities (tasks/actors) to use the acquired resources.
|
||||
|
||||
The first entity will be associated with the first bundle, the second
|
||||
entity will be associated with the second bundle, etc.
|
||||
|
||||
Args:
|
||||
entities: Remote Ray entities to annotate with the acquired resources.
|
||||
|
||||
Returns:
|
||||
The list of annotated remote Ray entities.
|
||||
"""
|
||||
bundles = self.resource_request.bundles
|
||||
|
||||
# Also count the empty head bundle as a bundle
|
||||
num_bundles = len(bundles) + int(self.resource_request.head_bundle_is_empty)
|
||||
|
||||
if len(entities) > num_bundles:
|
||||
raise RuntimeError(
|
||||
f"The number of callables to annotate ({len(entities)}) cannot "
|
||||
f"exceed the number of available bundles ({num_bundles})."
|
||||
)
|
||||
|
||||
annotated = []
|
||||
|
||||
if self.resource_request.head_bundle_is_empty:
|
||||
# The empty head bundle is place on the first bundle index with empty
|
||||
# resources.
|
||||
annotated.append(
|
||||
self._annotate_remote_entity(entities[0], {}, bundle_index=0)
|
||||
)
|
||||
|
||||
# Shift the remaining entities
|
||||
entities = entities[1:]
|
||||
|
||||
for i, (entity, bundle) in enumerate(zip(entities, bundles)):
|
||||
annotated.append(
|
||||
self._annotate_remote_entity(entity, bundle, bundle_index=i)
|
||||
)
|
||||
|
||||
return annotated
|
||||
|
||||
def _annotate_remote_entity(
|
||||
self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
|
||||
) -> RemoteRayEntity:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,155 @@
|
||||
import abc
|
||||
from typing import List, Optional
|
||||
|
||||
import ray
|
||||
from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ResourceManager(abc.ABC):
|
||||
"""Resource manager interface.
|
||||
|
||||
A resource manager can be used to request resources from a Ray cluster and
|
||||
allocate them to remote Ray tasks or actors.
|
||||
|
||||
Resources have to be requested before they can be acquired.
|
||||
|
||||
Resources managed by the resource manager can be in three states:
|
||||
|
||||
1. "Requested": The resources have been requested but are not yet available to
|
||||
schedule remote Ray objects. The resource request may trigger autoscaling,
|
||||
and can be cancelled if no longer needed.
|
||||
2. "Ready": The requested resources are now available to schedule remote Ray
|
||||
objects. They can be acquired and subsequently used remote Ray objects.
|
||||
The resource request can still be cancelled if no longer needed.
|
||||
3. "Acquired": The resources have been acquired by a caller to use for scheduling
|
||||
remote Ray objects. Note that it is the responsibility of the caller to
|
||||
schedule the Ray objects with these resources.
|
||||
The associated resource request has been completed and can no longer be
|
||||
cancelled. The acquired resources can be freed by the resource manager when
|
||||
they are no longer used.
|
||||
|
||||
The flow is as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Create resource manager
|
||||
resource_manager = ResourceManager()
|
||||
|
||||
# Create resource request
|
||||
resource_request = ResourceRequest([{"CPU": 4}])
|
||||
|
||||
# Pass to resource manager
|
||||
resource_manager.request_resources(resource_request)
|
||||
|
||||
# Wait until ready
|
||||
while not resource_manager.has_resources_ready(resource_request):
|
||||
time.sleep(1)
|
||||
|
||||
# Once ready, acquire resources
|
||||
acquired_resource = resource_manager.acquire_resources(resource_request)
|
||||
|
||||
# Bind to remote task or actor
|
||||
annotated_remote_fn = acquired_resource.annotate_remote_entities(
|
||||
[remote_fn])
|
||||
|
||||
# Run remote function. This will use the acquired resources
|
||||
ray.get(annotated_remote_fn.remote())
|
||||
|
||||
# After using the resources, free
|
||||
resource_manager.free_resources(annotated_resources)
|
||||
|
||||
"""
|
||||
|
||||
def request_resources(self, resource_request: ResourceRequest):
|
||||
"""Request resources.
|
||||
|
||||
Depending on the backend, resources can trigger autoscaling. Requested
|
||||
resources can be ready or not ready. Once they are "ready", they can
|
||||
be acquired and used by remote Ray objects.
|
||||
|
||||
Resource requests can be cancelled anytime using ``cancel_resource_request()``.
|
||||
Once acquired, the resource request is removed. Acquired resources can be
|
||||
freed with ``free_resources()``.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def cancel_resource_request(self, resource_request: ResourceRequest):
|
||||
"""Cancel resource request.
|
||||
|
||||
Resource requests can be cancelled anytime before a resource is acquired.
|
||||
Acquiring a resource will remove the associated resource request.
|
||||
Acquired resources can be freed with ``free_resources()``.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
|
||||
"""Returns True if resources for the given request are ready to be acquired."""
|
||||
raise NotImplementedError
|
||||
|
||||
def acquire_resources(
|
||||
self, resource_request: ResourceRequest
|
||||
) -> Optional[AcquiredResources]:
|
||||
"""Acquire resources. Returns None if resources are not ready to be acquired.
|
||||
|
||||
Acquiring resources will remove the associated resource request.
|
||||
Acquired resources can be returned with ``free_resources()``.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def free_resources(self, acquired_resource: AcquiredResources):
|
||||
"""Free acquired resources from usage and return them to the resource manager.
|
||||
|
||||
Freeing resources will return the resources to the manager, but there are
|
||||
no guarantees about the tasks and actors scheduled on the resources. The caller
|
||||
should make sure that any references to tasks or actors scheduled on the
|
||||
resources have been removed before calling ``free_resources()``.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_resource_futures(self) -> List[ray.ObjectRef]:
|
||||
"""Return futures for resources to await.
|
||||
|
||||
Depending on the backend, we use resource futures to determine availability
|
||||
of resources (e.g. placement groups) or resolution of requests.
|
||||
In this case, the futures can be awaited externally by the caller.
|
||||
|
||||
When a resource future resolved, the caller may call ``update_state()``
|
||||
to force the resource manager to update its internal state immediately.
|
||||
"""
|
||||
return []
|
||||
|
||||
def update_state(self):
|
||||
"""Update internal state of the resource manager.
|
||||
|
||||
The resource manager may have internal state that needs periodic updating.
|
||||
For instance, depending on the backend, resource futures can be awaited
|
||||
externally (with ``get_resource_futures()``).
|
||||
|
||||
If such a future resolved, the caller can instruct the resource
|
||||
manager to update its internal state immediately.
|
||||
"""
|
||||
pass
|
||||
|
||||
def clear(self):
|
||||
"""Reset internal state and clear all resources.
|
||||
|
||||
Calling this method will reset the resource manager to its initialization state.
|
||||
All resources will be removed.
|
||||
|
||||
Clearing the state will remove tracked resources from the manager, but there are
|
||||
no guarantees about the tasks and actors scheduled on the resources. The caller
|
||||
should make sure that any references to tasks or actors scheduled on the
|
||||
resources have been removed before calling ``clear()``.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __reduce__(self):
|
||||
"""We disallow serialization.
|
||||
|
||||
Shared resource managers should live on an actor.
|
||||
"""
|
||||
raise ValueError(
|
||||
f"Resource managers cannot be serialized. Resource manager: {str(self)}"
|
||||
)
|
||||
Reference in New Issue
Block a user