Ray Tasks

Use Ray tasks to parallelize Python functions across CPUs, GPUs, and nodes.

A Ray task is a stateless unit of computation. Annotate a function with @ray.remote and invoke it with .remote(...) to schedule it on the cluster.

Define and run tasks

import ray

ray.init()

@ray.remote
def square(x: int) -> int:
    return x * x

future = square.remote(4)
print(ray.get(future))  # 16

ray.get blocks on the result. To process many futures concurrently, collect them into a list and call ray.get once.

futures = [square.remote(i) for i in range(10)]
print(ray.get(futures))

Specify resource requirements

Declare CPU, GPU, memory, and custom resources on the decorator or per-call.

Decorator

@ray.remote(num_cpus=2, num_gpus=1)
def train(x):
    ...

Pass data to tasks

Tasks receive arguments by value. Large objects should be placed in the object store with ray.put so they aren’t re-serialized for every call.

import numpy as np

data = np.ones((10_000, 10_000))
ref = ray.put(data)

@ray.remote
def partial_sum(arr, start, end):
    return arr[start:end].sum()

futures = [partial_sum.remote(ref, i * 1000, (i + 1) * 1000) for i in range(10)]
print(sum(ray.get(futures)))

You can also pass ObjectRef values directly between tasks; Ray will resolve them on the destination worker.

Multiple return values

Use num_returns to return multiple objects.

@ray.remote(num_returns=2)
def mean_and_std(xs):
    return sum(xs) / len(xs), (sum((x - sum(xs)/len(xs))**2 for x in xs) / len(xs))**0.5

mean_ref, std_ref = mean_and_std.remote([1, 2, 3, 4])

Wait for results

Use ray.wait to process the first N results as they finish.

refs = [square.remote(i) for i in range(100)]
done, pending = ray.wait(refs, num_returns=10)

Generators and streaming returns

For long-running tasks that produce data incrementally, use a remote generator.

@ray.remote
def stream():
    for i in range(10):
        yield i

for ref in stream.remote():
    print(ray.get(ref))

Cancel tasks

ref = long_running.remote()
ray.cancel(ref)

Retries on failure

By default, Ray retries failed tasks up to three times. Configure with max_retries.

@ray.remote(max_retries=5)
def flaky():
    ...

Use retry_exceptions=True to retry on application-level exceptions, not just worker failures.

For long pipelines, prefer many small tasks over one huge task. The scheduler can pack small tasks more efficiently and recover from failures with less wasted work.

Next steps

Actors

Persistent stateful workers.

Patterns

Effective ways to structure tasks in real applications.