chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,385 @@
|
||||
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."
|
||||
)
|
||||
Reference in New Issue
Block a user