wehub-resource-sync
257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
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,
|
|
)
|