wehub-resource-sync
386 lines
14 KiB
Python
386 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import collections
|
|
from typing import TYPE_CHECKING, Deque, Iterator, Optional
|
|
|
|
import ray
|
|
from ray.exceptions import GetTimeoutError, ObjectRefStreamEndOfStreamError
|
|
from ray.util.annotations import DeveloperAPI, PublicAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray._private.worker import Worker
|
|
|
|
|
|
@DeveloperAPI
|
|
class DynamicObjectRefGenerator:
|
|
def __init__(self, refs: Deque["ray.ObjectRef"]):
|
|
# TODO(swang): As an optimization, can also store the generator
|
|
# ObjectID so that we don't need to keep individual ref counts for the
|
|
# inner ObjectRefs.
|
|
self._refs: Deque["ray.ObjectRef"] = collections.deque(refs)
|
|
|
|
def __iter__(self) -> Iterator("ray.ObjectRef"):
|
|
while self._refs:
|
|
yield self._refs.popleft()
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._refs)
|
|
|
|
|
|
@PublicAPI
|
|
class ObjectRefGenerator:
|
|
"""A generator to obtain object references from a task in a streaming manner.
|
|
|
|
The class is compatible with the Python generator and async generator interfaces.
|
|
|
|
The class is not thread-safe.
|
|
|
|
Do not initialize the class and create an instance directly.
|
|
The instance should be created by `.remote`.
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from typing import Generator
|
|
|
|
@ray.remote(num_returns="streaming")
|
|
def gen() -> Generator[int, None, None]:
|
|
for i in range(5):
|
|
yield i
|
|
|
|
obj_ref_gen: ray.ObjectRefGenerator = gen.remote()
|
|
for obj_ref in obj_ref_gen:
|
|
print("Got:", ray.get(obj_ref))
|
|
"""
|
|
|
|
def __init__(self, generator_ref: "ray.ObjectRef", worker: "Worker"):
|
|
# The reference to a generator task.
|
|
self._generator_ref = generator_ref
|
|
# True if an exception has been raised from the generator task.
|
|
self._generator_task_raised = False
|
|
# Ray's worker class. ray._private.worker.global_worker
|
|
self.worker = worker
|
|
self.worker.check_connected()
|
|
assert hasattr(worker, "core_worker")
|
|
|
|
# Public APIs
|
|
|
|
def __iter__(self) -> "ObjectRefGenerator":
|
|
return self
|
|
|
|
def __next__(self) -> "ray.ObjectRef":
|
|
"""Waits until a next ref is available and returns the object ref.
|
|
|
|
Raises StopIteration if there's no more objects
|
|
to generate.
|
|
|
|
The object ref will contain an exception if the task fails.
|
|
When the generator task returns N objects, it can return
|
|
up to N + 1 objects (if there's a system failure, the
|
|
last object will contain a system level exception).
|
|
"""
|
|
return self._next_sync()
|
|
|
|
def send(self, value):
|
|
raise NotImplementedError("`gen.send` is not supported.")
|
|
|
|
def throw(self, value):
|
|
raise NotImplementedError("`gen.throw` is not supported.")
|
|
|
|
def close(self):
|
|
raise NotImplementedError("`gen.close` is not supported.")
|
|
|
|
def __aiter__(self) -> "ObjectRefGenerator":
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
return await self._next_async()
|
|
|
|
async def asend(self, value):
|
|
raise NotImplementedError("`gen.asend` is not supported.")
|
|
|
|
async def athrow(self, value):
|
|
raise NotImplementedError("`gen.athrow` is not supported.")
|
|
|
|
async def aclose(self):
|
|
raise NotImplementedError("`gen.aclose` is not supported.")
|
|
|
|
def completed(self) -> "ray.ObjectRef":
|
|
"""Returns an object ref that is ready when
|
|
a generator task completes.
|
|
|
|
If the task is failed unexpectedly (e.g., worker failure),
|
|
the `ray.get(gen.completed())` raises an exception.
|
|
|
|
The function returns immediately.
|
|
"""
|
|
return self._generator_ref
|
|
|
|
def next_ready(self) -> bool:
|
|
"""If True, it means the output of next(gen) is ready and
|
|
ray.get(next(gen)) returns immediately. False otherwise.
|
|
|
|
It returns False when next(gen) raises a StopIteration
|
|
(this condition should be checked using is_finished).
|
|
|
|
The function returns immediately.
|
|
"""
|
|
self.worker.check_connected()
|
|
core_worker = self.worker.core_worker
|
|
|
|
if self.is_finished():
|
|
return False
|
|
|
|
expected_ref, is_ready = core_worker.peek_object_ref_stream(self._generator_ref)
|
|
|
|
if is_ready:
|
|
return True
|
|
|
|
ready, _ = ray.wait([expected_ref], timeout=0, fetch_local=False)
|
|
return len(ready) > 0
|
|
|
|
def is_finished(self) -> bool:
|
|
"""If True, it means the generator is finished
|
|
and all output is taken. False otherwise.
|
|
|
|
When True, if next(gen) is called, it will raise StopIteration
|
|
or StopAsyncIteration
|
|
|
|
The function returns immediately.
|
|
"""
|
|
self.worker.check_connected()
|
|
core_worker = self.worker.core_worker
|
|
|
|
finished = core_worker.is_object_ref_stream_finished(self._generator_ref)
|
|
|
|
if finished:
|
|
if self._generator_task_raised:
|
|
return True
|
|
else:
|
|
# We should try ray.get on a generator ref.
|
|
# If it raises an exception and
|
|
# _generator_task_raised is not set,
|
|
# this means the last ref is not taken yet.
|
|
try:
|
|
ray.get(self._generator_ref)
|
|
except Exception:
|
|
# The exception from _generator_ref
|
|
# hasn't been taken yet.
|
|
return False
|
|
else:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
# Private APIs
|
|
|
|
def _get_next_object_id_binary(self) -> bytes:
|
|
"""Return the binary id of the next object in the stream."""
|
|
self.worker.check_connected()
|
|
return self.worker.core_worker.peek_next_object_id_binary(self._generator_ref)
|
|
|
|
def _stream_exhausted(self) -> bool:
|
|
"""Whether the stream's end-of-stream marker has been reached and all
|
|
yielded refs consumed.
|
|
|
|
Non-blocking, in-memory check (unlike ``is_finished``, this does not
|
|
``ray.get`` the generator return object). When True, the only thing
|
|
left is the end-of-stream ``ray.get`` of the return object that
|
|
``_next_sync`` performs to surface ``StopIteration`` / task errors.
|
|
"""
|
|
self.worker.check_connected()
|
|
return self.worker.core_worker.is_object_ref_stream_finished(
|
|
self._generator_ref
|
|
)
|
|
|
|
def _get_next_ref_n(self, num_refs: int) -> list["ray.ObjectRef"]:
|
|
"""Return the next num_refs references from a generator without consuming them.
|
|
|
|
The returned refs are not consumed; wait for the last one to become ready
|
|
before calling ``_consume_next_ref_n`` to advance the stream.
|
|
|
|
Args:
|
|
num_refs: The number of references to return, starting from the
|
|
current head of the stream. Must be positive.
|
|
|
|
Returns:
|
|
A list of exactly num_refs ObjectRefs corresponding to the next
|
|
results in the stream, starting from the current head.
|
|
"""
|
|
if num_refs <= 0:
|
|
raise ValueError("num_refs must be positive")
|
|
self.worker.check_connected()
|
|
core_worker = self.worker.core_worker
|
|
return [
|
|
ref
|
|
for ref, _ in core_worker.peek_object_ref_stream_n(
|
|
self._generator_ref, num_refs
|
|
)
|
|
]
|
|
|
|
def _consume_next_ref_n(self, num_refs: int) -> None:
|
|
"""Consume (advance) the next num_refs references from a generator.
|
|
|
|
The caller must have waited for the last requested ref to become ready
|
|
(see ``_get_next_ref_n``); otherwise this raises ``ValueError`` instead
|
|
of silently advancing past unwritten objects.
|
|
|
|
If fewer than num_refs references remain before the end of the stream,
|
|
only the remaining references are consumed and the call returns
|
|
without raising.
|
|
|
|
Args:
|
|
num_refs: The number of references to consume, starting from the
|
|
current head of the stream. Must be positive.
|
|
"""
|
|
if num_refs <= 0:
|
|
raise ValueError("num_refs must be positive")
|
|
self.worker.check_connected()
|
|
core_worker = self.worker.core_worker
|
|
try:
|
|
core_worker.try_read_next_object_ref_stream_n(self._generator_ref, num_refs)
|
|
except ObjectRefStreamEndOfStreamError:
|
|
return
|
|
|
|
def _next_sync(self, timeout_s: Optional[int | float] = None) -> "ray.ObjectRef":
|
|
"""Waits for timeout_s and returns the object ref if available.
|
|
|
|
If an object is not available within the given timeout, it
|
|
returns a nil object reference.
|
|
|
|
If -1 timeout is provided, it means it waits infinitely.
|
|
|
|
Waiting is implemented as busy waiting.
|
|
|
|
Raises StopIteration if there's no more objects
|
|
to generate.
|
|
|
|
The object ref will contain an exception if the task fails.
|
|
When the generator task returns N objects, it can return
|
|
up to N + 1 objects (if there's a system failure, the
|
|
last object will contain a system level exception).
|
|
|
|
Args:
|
|
timeout_s: If the next object is not ready within
|
|
this timeout, it returns the nil object ref.
|
|
|
|
Returns:
|
|
ObjectRef corresponding to the next result in the stream.
|
|
"""
|
|
core_worker = self.worker.core_worker
|
|
|
|
# Wait for the next ObjectRef to become ready.
|
|
expected_ref, is_ready = core_worker.peek_object_ref_stream(self._generator_ref)
|
|
|
|
if not is_ready:
|
|
_, unready = ray.wait([expected_ref], timeout=timeout_s, fetch_local=False)
|
|
if len(unready) > 0:
|
|
return ray.ObjectRef.nil()
|
|
|
|
try:
|
|
ref = core_worker.try_read_next_object_ref_stream(self._generator_ref)
|
|
assert not ref.is_nil()
|
|
except ObjectRefStreamEndOfStreamError:
|
|
if self._generator_task_raised:
|
|
# Exception has been returned.
|
|
raise StopIteration from None
|
|
|
|
try:
|
|
# The generator ref contains an exception
|
|
# if there's any failure. It contains nothing otherwise.
|
|
# In that case, it should raise StopIteration.
|
|
#
|
|
# Bound this get by the caller's timeout: the return object
|
|
# can be remote — or lost to a failed node and pending
|
|
# reconstruction — and an unbounded get would block the
|
|
# caller until it is restored (e.g. the Ray Data scheduling
|
|
# thread; a saturated cluster can then deadlock, since the
|
|
# blocked consumer is what releases backpressured CPUs).
|
|
# Per this method's contract, a timeout is reported as "no
|
|
# object ready yet" (nil ref) so the caller retries.
|
|
ray.get(
|
|
self._generator_ref,
|
|
timeout=(None if timeout_s is None or timeout_s < 0 else timeout_s),
|
|
)
|
|
except GetTimeoutError:
|
|
return ray.ObjectRef.nil()
|
|
except Exception:
|
|
self._generator_task_raised = True
|
|
return self._generator_ref
|
|
else:
|
|
# The task finished without an exception.
|
|
raise StopIteration from None
|
|
return ref
|
|
|
|
async def _suppress_exceptions(self, ref: "ray.ObjectRef") -> None:
|
|
# Wrap a streamed ref to avoid asyncio warnings about not retrieving
|
|
# the exception when we are just waiting for the ref to become ready.
|
|
# The exception will get returned (or warned) to the user once they
|
|
# actually await the ref.
|
|
try:
|
|
await ref
|
|
except Exception:
|
|
pass
|
|
|
|
async def _next_async(self, timeout_s: Optional[int | float] = None):
|
|
"""Same API as _next_sync, but it is for async context."""
|
|
core_worker = self.worker.core_worker
|
|
ref, is_ready = core_worker.peek_object_ref_stream(self._generator_ref)
|
|
|
|
if not is_ready:
|
|
# TODO(swang): Avoid fetching the value.
|
|
_, unready = await asyncio.wait(
|
|
[asyncio.create_task(self._suppress_exceptions(ref))], timeout=timeout_s
|
|
)
|
|
if len(unready) > 0:
|
|
return ray.ObjectRef.nil()
|
|
|
|
try:
|
|
ref = core_worker.try_read_next_object_ref_stream(self._generator_ref)
|
|
assert not ref.is_nil()
|
|
except ObjectRefStreamEndOfStreamError:
|
|
if self._generator_task_raised:
|
|
# Exception has been returned.
|
|
raise StopAsyncIteration from None
|
|
|
|
try:
|
|
# The generator ref contains an exception
|
|
# if there's any failure. It contains nothing otherwise.
|
|
# In that case, it should raise StopAsyncIteration.
|
|
#
|
|
# Bound this await by the caller's timeout, mirroring
|
|
# _next_sync: the return object can be remote — or lost to a
|
|
# failed node and pending reconstruction — and an unbounded
|
|
# await would block the caller until it is restored. Per this
|
|
# method's contract, a timeout is reported as "no object
|
|
# ready yet" (nil ref) so the caller retries.
|
|
if timeout_s is None or timeout_s < 0:
|
|
await self._generator_ref
|
|
else:
|
|
await asyncio.wait_for(self._generator_ref, timeout=timeout_s)
|
|
except asyncio.TimeoutError:
|
|
return ray.ObjectRef.nil()
|
|
except Exception:
|
|
self._generator_task_raised = True
|
|
return self._generator_ref
|
|
else:
|
|
# Meaning the task succeed without failure raise StopAsyncIteration.
|
|
raise StopAsyncIteration from None
|
|
|
|
return ref
|
|
|
|
def __del__(self):
|
|
if hasattr(self.worker, "core_worker"):
|
|
# The stream is created when a task is first submitted.
|
|
# NOTE: This can be called multiple times
|
|
# because python doesn't guarantee __del__ is called
|
|
# only once.
|
|
self.worker.core_worker.async_delete_object_ref_stream(self._generator_ref)
|
|
|
|
def __getstate__(self):
|
|
raise TypeError(
|
|
"You cannot return or pass a generator to other task. "
|
|
"Serializing a ObjectRefGenerator is not allowed."
|
|
)
|