Pattern: Using ray.wait to limit the number of pending tasks

In this pattern, you use ray.wait() to limit the number of pending tasks.

If you continuously submit tasks faster than they finish processing, tasks accumulate in the pending task queue, which can eventually cause out-of-memory (OOM) errors. Use ray.wait() to apply backpressure and limit the number of pending tasks, so the pending task queue doesn’t grow indefinitely and cause OOM.

Note

If you submit a finite number of tasks, you’re unlikely to hit this issue, because each task uses only a small amount of memory for bookkeeping in the queue. The issue is more likely when you have an infinite stream of tasks to run.

Note

Use this method primarily to limit how many tasks are in flight at the same time. You can also use it to limit how many tasks run concurrently, but avoid doing so, because it can hurt scheduling performance. Ray decides task parallelism automatically based on resource availability, so to adjust how many tasks can run concurrently, modify each task’s resource requirements instead.

Example use case

You have a worker actor that processes tasks at a rate of X tasks per second, and you want to submit tasks to it at a rate lower than X to avoid OOM.

For example, Ray Serve uses this pattern to limit the number of pending queries for each worker.

Limit the number of pending tasks

Code example

Without backpressure:

import ray

ray.init()


@ray.remote
class Actor:
    async def heavy_compute(self):
        # taking a long time...
        # await asyncio.sleep(5)
        return


actor = Actor.remote()

NUM_TASKS = 1000
result_refs = []
# When NUM_TASKS is large enough, this will eventually OOM.
for _ in range(NUM_TASKS):
    result_refs.append(actor.heavy_compute.remote())
ray.get(result_refs)

With backpressure:

MAX_NUM_PENDING_TASKS = 100
result_refs = []
for _ in range(NUM_TASKS):
    if len(result_refs) > MAX_NUM_PENDING_TASKS:
        # update result_refs to only
        # track the remaining tasks.
        ready_refs, result_refs = ray.wait(result_refs, num_returns=1)
        ray.get(ready_refs)

    result_refs.append(actor.heavy_compute.remote())

ray.get(result_refs)