wehub-resource-sync
149 lines
4.8 KiB
Python
149 lines
4.8 KiB
Python
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)
|