Ray generators#

Choosing a generator API#

Ray’s streaming generator API, also called the regular generator API, is the recommended way to consume generator results. It returns an ObjectRefGenerator. Use it when a task yields a known or naturally bounded stream of results. You can iterate over it, call ray.get on each object reference, or pass it to ray.wait.

You need the deprecated dynamic generator API only to support existing code that sets num_returns="dynamic". It returns a DynamicObjectRefGenerator through a single ObjectRef. In new code, use the streaming generator API instead. Ray plans to remove the dynamic generator API in a future version. For details, see Dynamic generators.

Python generators are functions that behave like iterators, yielding one value per iteration. Ray also supports the Python generator API.

Any generator function decorated with ray.remote becomes a Ray generator task. Generator tasks stream outputs back to the caller before the task finishes.

+import ray
 import time

 # Takes 25 seconds to finish.
+@ray.remote
 def f():
     for i in range(5):
         time.sleep(5)
         yield i

-for obj in f():
+for obj_ref in f.remote():
     # Prints every 5 seconds and stops after 25 seconds.
-    print(obj)
+    print(ray.get(obj_ref))

The preceding Ray generator yields an output every 5 seconds, five times. With a normal Ray task, you have to wait 25 seconds to access the output. With a Ray generator, the caller can access the object reference before the task f finishes.

Use a Ray generator in the following cases:

You want to reduce heap memory or object store memory usage by yielding and garbage-collecting the output before the task finishes.
You’re familiar with Python generators and want an equivalent programming model.

Ray libraries use Ray generators to support streaming use cases:

Ray Serve uses Ray generators to support streaming responses.
Ray Data is a streaming data processing library that uses Ray generators to control and reduce concurrent memory usage.

Ray generators work with existing Ray APIs:

You can use Ray generators in both actor and non-actor tasks.
Ray generators work with all actor execution models, including threaded actors and async actors.
Ray generators work with built-in fault tolerance features such as retry or lineage reconstruction.
Ray generators work with Ray APIs such as ray.wait and ray.cancel.

Getting started#

Define a Python generator function and decorate it with ray.remote to create a Ray generator.

import ray
import time

@ray.remote
def task():
    for i in range(5):
        time.sleep(5)
        yield i

A Ray generator task returns an ObjectRefGenerator object, which is compatible with the generator and async generator APIs. You can use the next, __iter__, __anext__, and __aiter__ APIs from the class.

Each time a task invokes yield, the corresponding output becomes available from the generator as a Ray object reference. Call next(gen) to get an object reference. If next has no more items to generate, it raises StopIteration. If __anext__ has no more items to generate, it raises StopAsyncIteration.

The next API blocks the thread until the task generates the next object reference with yield. You can also iterate over object references with a for loop.

To avoid blocking a thread, use asyncio or the ray.wait API.

gen = task.remote()
# Blocks for 5 seconds.
ref = next(gen)
# return 0
ray.get(ref)
# Blocks for 5 seconds.
ref = next(gen)
# Return 1
ray.get(ref)

# Returns 2~4 every 5 seconds.
for ref in gen:
    print(ray.get(ref))

Note

A normal Python generator function pauses and resumes each time you call next on the generator. Ray eagerly executes a generator task to completion, regardless of whether the caller polls the partial results.

Error handling#

If a generator task fails because of an application exception or a system error, such as an unexpected node failure, next(gen) returns an object reference that contains an exception. When you call ray.get on that object reference, Ray raises the exception.

@ray.remote
def task():
    for i in range(5):
        time.sleep(1)
        if i == 1:
            raise ValueError
        yield i

gen = task.remote()
# it's okay.
ray.get(next(gen))

# Raises an exception
try:
    ray.get(next(gen))
except ValueError as e:
    print(f"Exception is raised when i == 1 as expected {e}")

In the preceding example, if the application fails the task, Ray returns the object reference with the exception in the correct order. For example, if Ray raises the exception after the second yield, the third next(gen) always returns an object reference with an exception. If a system error, such as a node failure or worker process failure, fails the task, next(gen) returns the object reference that contains the system-level exception at any time, without an ordering guarantee. As a result, when a generator has N yields and failures occur, it can create from 1 to N + 1 object references, up to N outputs and one object reference that contains the system-level exception.

Generator from actor tasks#

The Ray generator is compatible with all actor execution models. It works with regular actors, async actors, and threaded actors.

@ray.remote
class Actor:
    def f(self):
        for i in range(5):
            yield i

@ray.remote
class AsyncActor:
    async def f(self):
        for i in range(5):
            yield i

@ray.remote(max_concurrency=5)
class ThreadedActor:
    def f(self):
        for i in range(5):
            yield i

actor = Actor.remote()
for ref in actor.f.remote():
    print(ray.get(ref))

actor = AsyncActor.remote()
for ref in actor.f.remote():
    print(ray.get(ref))

actor = ThreadedActor.remote()
for ref in actor.f.remote():
    print(ray.get(ref))

Using the Ray generator with asyncio#

The returned ObjectRefGenerator is also compatible with asyncio. You can use __anext__ or async for loops.

import asyncio

@ray.remote
def task():
    for i in range(5):
        time.sleep(1)
        yield i


async def main():
    async for ref in task.remote():
        print(await ref)

asyncio.run(main())

Garbage collection of object references#

The object reference that next(generator) returns is a regular Ray object reference, and Ray applies the same distributed reference counting to it. If you don’t consume references from a generator with the next API, those references are garbage-collected when the generator is garbage-collected.

@ray.remote
def task():
    for i in range(5):
        time.sleep(1)
        yield i

gen = task.remote()
ref1 = next(gen)
del gen

In the preceding example, Ray counts ref1 as a normal Ray object reference after Ray returns it. Other references that you don’t consume with next(gen) are removed when the generator is garbage-collected. In this example, garbage collection happens when you call del gen.

Fault tolerance#

Fault tolerance features work with Ray generator tasks and actor tasks. For example, the following features work with Ray generators:

Task fault tolerance features: max_retries, retry_exceptions
Actor fault tolerance features: max_restarts, max_task_retries
Object fault tolerance features: object reconstruction

Cancellation#

The ray.cancel() function works with both Ray generator tasks and actor tasks. Semantically, canceling a generator task is the same as canceling a regular task. When you cancel a task, next(gen) can return an object reference that contains TaskCancelledError without any special ordering guarantee.

Wait for a generator without blocking a thread#

The next API blocks its thread until the next object reference is available. You can wait for a generator without blocking a thread in three ways.

Wait until a generator task completes#

ObjectRefGenerator has a completed API, which returns an object reference that becomes available when the generator task finishes or errors. For example, call ray.get(<generator_instance>.completed()) to wait until the task completes. You can’t pass an ObjectRefGenerator to ray.get directly.

Use asyncio and await#

ObjectRefGenerator is compatible with asyncio. To avoid blocking a thread, create multiple asyncio tasks that create a generator task and wait for it.

import asyncio

@ray.remote
def task():
    for i in range(5):
        time.sleep(1)
        yield i


async def async_task():
    async for ref in task.remote():
        print(await ref)

async def main():
    t1 = async_task()
    t2 = async_task()
    await asyncio.gather(t1, t2)

asyncio.run(main())

Use ray.wait#

You can pass an ObjectRefGenerator as an input to ray.wait. The generator is “ready” if its next item is available. Once ray.wait returns the generator in the ready list, next(gen) returns the next object reference immediately without blocking. The following example shows this pattern.

@ray.remote
def task():
    for i in range(5):
        time.sleep(5)
        yield i

gen = task.remote()

# Because it takes 5 seconds to make the first yield,
# with 0 timeout, the generator is unready.
ready, unready = ray.wait([gen], timeout=0)
print("timeout 0, nothing is ready.")
print(ready)
assert len(ready) == 0
assert len(unready) == 1

# Without a timeout argument, ray.wait waits until the given argument
# is ready. When a next item is ready, it returns.
ready, unready = ray.wait([gen])
print("Wait for 5 seconds. The next item is ready.")
assert len(ready) == 1
assert len(unready) == 0
next(gen)

# Because the second yield hasn't happened yet,
ready, unready = ray.wait([gen], timeout=0)
print("Wait for 0 seconds. The next item is not ready.")
print(ready, unready)
assert len(ready) == 0
assert len(unready) == 1

All of the ray.wait input arguments, such as timeout, num_returns, and fetch_local, work with a generator.

You can mix regular Ray object references and generators in the inputs to ray.wait. In this case, your application should handle the two kinds of inputs differently. The following example checks whether each ready input is an ObjectRefGenerator.

from ray._raylet import ObjectRefGenerator

@ray.remote
def generator_task():
    for i in range(5):
        time.sleep(5)
        yield i

@ray.remote
def regular_task():
    for i in range(5):
        time.sleep(5)
    return

gen = [generator_task.remote()]
ref = [regular_task.remote()]
ready, unready = [], [*gen, *ref]
result = []

while unready:
    ready, unready = ray.wait(unready)
    for r in ready:
        if isinstance(r, ObjectRefGenerator):
            try:
                ref = next(r)
                result.append(ray.get(ref))
            except StopIteration:
                pass
            else:
                unready.append(r)
        else:
            result.append(ray.get(r))

Thread safety#

ObjectRefGenerator objects aren’t thread-safe.

Limitation#

Ray generators don’t support these features:

throw, send, and close APIs
return statements from generators
Passing ObjectRefGenerator to another task or actor
Ray Client

Deprecated dynamic generator#

Dynamic generators