Dynamic generators#

Warning

num_returns="dynamic" generator API is deprecated and scheduled for removal in an upcoming version. Use the streaming generator API instead. Ray emits a runtime RayDeprecationWarning when you use num_returns="dynamic".

Python generators are functions that behave like iterators, yielding one value per iteration. Ray supports remote generators for two use cases:

  1. Reducing maximum heap memory usage when a remote function returns multiple values. See the design pattern guide for an example.

  2. Returning a number of values that the remote function sets dynamically, instead of a number that the caller sets.

You can use remote generators in both actor and non-actor tasks.

num_returns set by the task caller#

Where possible, set the remote function’s number of return values from the caller with @ray.remote(num_returns=x) or foo.options(num_returns=x).remote(). Ray returns this many ObjectRefs to the caller, and the remote task should return the same number of values, usually as a tuple or list. Compared to setting the number of return values dynamically, this approach adds less complexity to your code and less performance overhead, because Ray knows ahead of time exactly how many ObjectRefs to return to the caller.

Without changing the caller’s syntax, you can also use a remote generator function to yield the values iteratively. The generator should yield the same number of return values that the caller specifies, and Ray stores them one at a time in its object store. If a generator yields a different number of values than the caller specifies, Ray raises an error.

For example, you can replace the following code, which returns a list of return values:

import numpy as np


@ray.remote
def large_values(num_returns):
    return [
        np.random.randint(np.iinfo(np.int8).max, size=(100_000_000, 1), dtype=np.int8)
        for _ in range(num_returns)
    ]

with this code, which uses a generator function:

@ray.remote
def large_values_generator(num_returns):
    for i in range(num_returns):
        yield np.random.randint(
            np.iinfo(np.int8).max, size=(100_000_000, 1), dtype=np.int8
        )
        print(f"yielded return value {i}")

With this change, the generator function doesn’t need to hold all of its return values in memory at once. It can yield the arrays one at a time to reduce memory pressure.

num_returns set by the task executor#

In some cases, the caller doesn’t know how many return values to expect from a remote function. For example, suppose you want to write a task that breaks up its argument into equal-size chunks and returns them. You might not know the size of the argument until the task executes, so you don’t know how many return values to expect.

In these cases, you can use a remote generator function that returns a dynamic number of values. To use this feature, set num_returns="dynamic" in the @ray.remote decorator or the remote function’s .options(). When you invoke the remote function, Ray returns a single ObjectRef, which it populates with a DynamicObjectRefGenerator when the task completes. Use the DynamicObjectRefGenerator to iterate over a list of ObjectRefs that contain the actual values the task returned.

import numpy as np


@ray.remote(num_returns="dynamic")
def split(array, chunk_size):
    while len(array) > 0:
        yield array[:chunk_size]
        array = array[chunk_size:]


array_ref = ray.put(np.zeros(np.random.randint(1000_000)))
block_size = 1000

# Returns an ObjectRef[DynamicObjectRefGenerator].
dynamic_ref = split.remote(array_ref, block_size)
print(dynamic_ref)
# ObjectRef(c8ef45ccd0112571ffffffffffffffffffffffff0100000001000000)

i = -1
ref_generator = ray.get(dynamic_ref)
print(ref_generator)
# <ray._raylet.DynamicObjectRefGenerator object at 0x7f7e2116b290>
for i, ref in enumerate(ref_generator):
    # Each DynamicObjectRefGenerator iteration returns an ObjectRef.
    assert len(ray.get(ref)) <= block_size
num_blocks_generated = i + 1
array_size = len(ray.get(array_ref))
assert array_size <= num_blocks_generated * block_size
print(f"Split array of size {array_size} into {num_blocks_generated} blocks of "
      f"size {block_size} each.")
# Split array of size 63153 into 64 blocks of size 1000 each.

# NOTE: The dynamic_ref points to the generated ObjectRefs. Make sure that this
# ObjectRef goes out of scope so that Ray can garbage-collect the internal
# ObjectRefs.
del dynamic_ref

You can also pass the ObjectRef that a task with num_returns="dynamic" returns to another task. The receiving task gets the DynamicObjectRefGenerator, which it can use to iterate over the original task’s return values. You can also pass the DynamicObjectRefGenerator itself as a task argument, but avoid this pattern. Getting the generator takes an extra ray.get call, which blocks the driver.

@ray.remote
def get_size(ref_generator : DynamicObjectRefGenerator):
    print(ref_generator)
    num_elements = 0
    for ref in ref_generator:
        array = ray.get(ref)
        assert len(array) <= block_size
        num_elements += len(array)
    return num_elements


# Returns an ObjectRef[DynamicObjectRefGenerator].
dynamic_ref = split.remote(array_ref, block_size)
assert array_size == ray.get(get_size.remote(dynamic_ref))
# (get_size pid=1504184)
# <ray._raylet.DynamicObjectRefGenerator object at 0x7f81c4250ad0>

# This also works, but should be avoided because you have to call an additional
# `ray.get`, which blocks the driver.
ref_generator = ray.get(dynamic_ref)
assert array_size == ray.get(get_size.remote(ref_generator))
# (get_size pid=1504184)
# <ray._raylet.DynamicObjectRefGenerator object at 0x7f81c4251b50>

Exception handling#

If a generator function raises an exception before yielding all its values, the values that it already stored remain accessible through their ObjectRefs. The remaining ObjectRefs contain the raised exception. This behavior applies to both static and dynamic num_returns. If you call the task with num_returns="dynamic", Ray stores the exception as an additional final ObjectRef in the DynamicObjectRefGenerator.

@ray.remote
def generator():
    for i in range(2):
        yield i
    raise Exception("error")


ref1, ref2, ref3, ref4 = generator.options(num_returns=4).remote()
assert ray.get([ref1, ref2]) == [0, 1]
# All remaining ObjectRefs will contain the error.
try:
    ray.get([ref3, ref4])
except Exception as error:
    print(error)

dynamic_ref = generator.options(num_returns="dynamic").remote()
ref_generator = ray.get(dynamic_ref)
ref1, ref2, ref3 = ref_generator
assert ray.get([ref1, ref2]) == [0, 1]
# Generators with num_returns="dynamic" will store the exception in the final
# ObjectRef.
try:
    ray.get(ref3)
except Exception as error:
    print(error)

A known bug currently prevents Ray from propagating exceptions for generators that yield more values than expected. This bug can occur in two cases:

  1. When the caller sets num_returns, but the generator task returns more values than that number.

  2. When Ray re-executes a generator task with num_returns="dynamic", and the re-executed task yields more values than the original execution.

In general, Ray doesn’t guarantee correctness for re-executing a nondeterministic task, so set @ray.remote(max_retries=0) for such tasks.

# Generators that yield more values than expected currently do not throw an
# exception (the error is only logged).
# See https://github.com/ray-project/ray/issues/28689.
ref1, ref2 = generator.options(num_returns=2).remote()
assert ray.get([ref1, ref2]) == [0, 1]
"""
(generator pid=2375938) 2022-09-28 11:08:51,386 ERROR worker.py:755 --
    Unhandled error: Task threw exception, but all return values already
    created.  This should only occur when using generator tasks.
...
"""

Limitations#

Although a generator function creates ObjectRefs one at a time, Ray currently doesn’t schedule dependent tasks until the entire task completes and has created all its values. This behavior is similar to the semantics of tasks that return multiple values as a list.