chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user