from typing import TYPE_CHECKING, Any, Iterable, Iterator, Optional, Sequence, Union import ray from ray.util.annotations import PublicAPI if TYPE_CHECKING: from ray import ObjectRef from ray.remote_function import RemoteFunction # ray.wait() has a default num_returns of 1. # Using a slightly larger batch until the optimization is fully implemented, see # https://github.com/ray-project/ray/issues/49905 DEFAULT_CHUNK_SIZE = 10 DEFAULT_BACKPRESSURE_SIZE = 100 def _wait_and_get_single_batch( refs: "Sequence[ObjectRef]", *, chunk_size: int, yield_obj_refs: bool = False, **kwargs, ) -> tuple[list[Union[Any, "ObjectRef"]], "list[ObjectRef]"]: """Call ray.wait and explicitly return the ready objects/results and remaining Ray remote refs. Args: refs: A list of Ray object refs. chunk_size: The `num_returns` parameter to pass to `ray.wait()`. yield_obj_refs: If True, return Ray remote refs instead of results (by calling :meth:`~ray.get`). **kwargs: Additional keyword arguments to pass to `ray.wait()`. Returns: A tuple of two lists, ready and not ready. This is the same as the return value of `ray.wait()`. """ if chunk_size < 1: raise ValueError("`chunk_size` must be >= 1") kwargs = kwargs or {} # num_returns must be <= len(refs) ready, refs = ray.wait( refs, num_returns=min(chunk_size, len(refs)), **kwargs, ) if not yield_obj_refs: return ray.get(ready), refs return ready, refs @PublicAPI(stability="alpha") def as_completed( refs: "Sequence[ObjectRef]", *, chunk_size: int = DEFAULT_CHUNK_SIZE, yield_obj_refs: bool = False, **kwargs, ) -> Iterator[Union[Any, "ObjectRef"]]: """Given a list of Ray task references, yield results as they become available. Unlike calling :meth:`~ray.get` on a list of references (i.e., `ray.get(refs)`) which waits for all results to be ready, this function begins to yield result as soon as a batch of `chunk_size` results are ready. .. note:: Generally there is no guarantee on the order of results. For example, the first result is not necessarily the first one completed, but rather the first one submitted in the first available batch (See :meth:`~ray.wait` for more details about preservation of submission order). .. note:: Use this function instead of calling :meth:`~ray.get` inside a for loop. See https://docs.ray.io/en/latest/ray-core/patterns/ray-get-loop.html for more details. Example: Suppose we have a function that sleeps for x seconds depending on the input. We expect to obtain a partially sorted list of results. .. testcode:: python import ray import time @ray.remote def f(x): time.sleep(x) return x refs = [f.remote(i) for i in [10, 4, 6, 8, 2]] for x in ray.util.as_completed(refs, chunk_size=2): print(x) .. testoutput:: :options: +MOCK # Output: 4 2 6 8 10 Args: refs: A list of Ray object refs. chunk_size: The number of tasks to wait for in each iteration (default 10). The parameter is passed as `num_returns` to :meth:`~ray.wait` internally. yield_obj_refs: If True, return Ray remote refs instead of results (by calling :meth:`~ray.get`). **kwargs: Additional keyword arguments to pass to :meth:`~ray.wait`, e.g., `timeout` and `fetch_local`. Yields: Union[Any, ObjectRef]: The results (or optionally their Ray references) of the Ray tasks as they complete. """ if chunk_size < 1: raise ValueError("`chunk_size` must be >= 1") if "num_returns" in kwargs: raise ValueError("Use the `chunksize` argument instead of `num_returns`.") while refs: results, refs = _wait_and_get_single_batch( refs, chunk_size=chunk_size, yield_obj_refs=yield_obj_refs, **kwargs, ) yield from results @PublicAPI(stability="alpha") def map_unordered( fn: "RemoteFunction", items: Iterable[Any], *, backpressure_size: Optional[int] = DEFAULT_BACKPRESSURE_SIZE, chunk_size: int = DEFAULT_CHUNK_SIZE, yield_obj_refs: bool = False, **kwargs, ) -> Iterator[Union[Any, "ObjectRef"]]: """Apply a Ray remote function to a list of items and return an iterator that yields the completed results as they become available. This helper function applies backpressure to control the number of pending tasks, following the design pattern described in https://docs.ray.io/en/latest/ray-core/patterns/limit-pending-tasks.html. .. note:: There is generally no guarantee on the order of results. Example: Suppose we have a function that sleeps for x seconds depending on the input. We expect to obtain a partially sorted list of results. .. testcode:: python import ray import time @ray.remote def f(x): time.sleep(x) return x # Example 1: chunk_size=2 for x in ray.util.map_unordered(f, [10, 4, 6, 8, 2], chunk_size=2): print(x) .. testoutput:: :options: +MOCK 4 2 6 8 10 .. testcode:: python # Example 2: backpressure_size=2, chunk_size=1 for x in ray.util.map_unordered(f, [10, 4, 6, 8, 2], backpressure_size=2, chunk_size=1): print(x) .. testoutput:: :options: +MOCK 4 10 6 8 2 Args: fn: A remote function to apply to the list of items. For more complex use cases, use Ray Data's :meth:`~ray.data.Dataset.map` / :meth:`~ray.data.Dataset.map_batches` instead. items: An iterable of items to apply the function to. backpressure_size: Maximum number of in-flight tasks allowed before calling a blocking :meth:`~ray.wait` (default 100). If None, no backpressure is applied. chunk_size: The number of tasks to wait for when the number of in-flight tasks exceeds `backpressure_size`. The parameter is passed as `num_returns` to :meth:`~ray.wait` internally. yield_obj_refs: If True, return Ray remote refs instead of results (by calling :meth:`~ray.get`). **kwargs: Additional keyword arguments to pass to :meth:`~ray.wait`, e.g., `timeout` and `fetch_local`. Yields: Union[Any, ObjectRef]: The results (or optionally their Ray references) of the Ray tasks as they complete. .. seealso:: :meth:`~ray.util.as_completed` Call this method for an existing list of Ray object refs. :meth:`~ray.data.Dataset.map` Use Ray Data APIs (e.g., :meth:`~ray.data.Dataset.map` and :meth:`~ray.data.Dataset.map_batches`) for better control and complex use cases, e.g., functions with multiple arguments. .. note:: This is an altenative to `pool.imap_unordered()` in Ray's Actor-based `multiprocessing.Pool`. See https://docs.ray.io/en/latest/ray-more-libs/multiprocessing.html for more details. """ if backpressure_size is None: backpressure_size: float = float("inf") elif backpressure_size <= 0: raise ValueError("backpressure_size must be positive.") if chunk_size < 1: raise ValueError("`chunk_size` must be >= 1") if "num_returns" in kwargs: raise ValueError("Use the `chunk_size` argument instead of `num_returns`.") refs = [] for item in items: refs.append(fn.remote(item)) if len(refs) >= backpressure_size: results, refs = _wait_and_get_single_batch( refs, chunk_size=chunk_size, yield_obj_refs=yield_obj_refs, **kwargs, ) yield from results else: yield from as_completed( refs, chunk_size=chunk_size, yield_obj_refs=yield_obj_refs, **kwargs, )