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." )