Lazy computation graphs with the Ray DAG API#

With ray.remote, your application runs its computation remotely at runtime. On a class or function decorated with ray.remote, you can also call .bind to build a static computation graph.

Note

Ray DAG is a developer-facing API with two recommended use cases:

Iterate locally on and test applications that you author with higher-level libraries.
Build libraries on top of the Ray DAG API.

Note

Ray has introduced an experimental API for high-performance workloads that is especially well suited for applications using multiple GPUs. This API is built on top of the Ray DAG API.

See Ray Compiled Graph for more details.

Calling .bind() on a class or function decorated with ray.remote generates an intermediate representation (IR) node. IR nodes are the backbone and building blocks of the DAG, and they statically hold the computation graph together. At execution time, each IR node resolves to a value in topological order.

You can also assign an IR node to a variable and pass it to other nodes as an argument.

Ray DAG with functions#

When you execute it, an IR node that .bind() generates from a ray.remote decorated function runs as a Ray task and resolves to the task’s output.

The following example builds a chain of functions. While you iterate, you can execute each node as the root node, or pass it to other functions as a positional or keyword argument to form more complex DAGs.

To execute any IR node directly as the root of the DAG, call dag_node.execute(). Execution ignores every node that isn’t reachable from the root.

Python

import ray

ray.init()

@ray.remote
def func(src, inc=1):
    return src + inc

a_ref = func.bind(1, inc=2)
assert ray.get(a_ref.execute()) == 3 # 1 + 2 = 3
b_ref = func.bind(a_ref, inc=3)
assert ray.get(b_ref.execute()) == 6 # (1 + 2) + 3 = 6
c_ref = func.bind(b_ref, inc=a_ref)
assert ray.get(c_ref.execute()) == 9 # ((1 + 2) + 3) + (1 + 2) = 9

Ray DAG with classes and class methods#

When you execute it, an IR node that .bind() generates from a ray.remote decorated class runs as a Ray actor. Ray instantiates the actor every time you execute the node. Calls to the class methods can form a chain of function calls specific to the parent actor instance.

You can combine IR nodes generated from functions, classes, and class methods to form a DAG.

Python

import ray

ray.init()

@ray.remote
class Actor:
    def __init__(self, init_value):
        self.i = init_value

    def inc(self, x):
        self.i += x

    def get(self):
        return self.i

a1 = Actor.bind(10)  # Instantiate Actor with init_value 10.
val = a1.get.bind()  # ClassMethod that returns value from get() from
                     # the actor created.
assert ray.get(val.execute()) == 10

@ray.remote
def combine(x, y):
    return x + y

a2 = Actor.bind(10) # Instantiate another Actor with init_value 10.
a1.inc.bind(2)  # Call inc() on the actor created with increment of 2.
a1.inc.bind(4)  # Call inc() on the actor created with increment of 4.
a2.inc.bind(6)  # Call inc() on the actor created with increment of 6.

# Combine outputs from a1.get() and a2.get()
dag = combine.bind(a1.get.bind(), a2.get.bind())

# a1 +  a2 + inc(2) + inc(4) + inc(6)
# 10 + (10 + ( 2   +    4    +   6)) = 32
assert ray.get(dag.execute()) == 32

Ray DAG with custom InputNode#

InputNode is the singleton node of a DAG that represents the input value at runtime. Use it as a context manager with no arguments, and supply its value as the arguments of dag_node.execute().

Python

import ray

ray.init()

from ray.dag.input_node import InputNode

@ray.remote
def a(user_input):
    return user_input * 2

@ray.remote
def b(user_input):
    return user_input + 1

@ray.remote
def c(x, y):
    return x + y

with InputNode() as dag_input:
    a_ref = a.bind(dag_input)
    b_ref = b.bind(dag_input)
    dag = c.bind(a_ref, b_ref)

#   a(2)  +   b(2)  = c
# (2 * 2) + (2 + 1)
assert ray.get(dag.execute(2)) == 7

#   a(3)  +   b(3)  = c
# (3 * 2) + (3 + 1)
assert ray.get(dag.execute(3)) == 10

Ray DAG with MultiOutputNode#

Use MultiOutputNode when a DAG has more than one output. dag_node.execute() returns a list of the object refs passed to MultiOutputNode. The following example shows a MultiOutputNode with two outputs.

Python

import ray

from ray.dag.input_node import InputNode
from ray.dag.output_node import MultiOutputNode

@ray.remote
def f(input):
    return input + 1

with InputNode() as input_data:
    dag = MultiOutputNode([f.bind(input_data["x"]), f.bind(input_data["y"])])

refs = dag.execute({"x": 1, "y": 2})
assert ray.get(refs) == [2, 3]

Reuse Ray actors in DAGs#

You can make actors part of the DAG definition with the Actor.bind() API. However, when a DAG finishes execution, Ray kills actors created with bind.

To keep your actors alive after the DAG finishes, create them with Actor.remote() instead.

Python

import ray
from ray.dag.input_node import InputNode
from ray.dag.output_node import MultiOutputNode

@ray.remote
class Worker:
    def __init__(self):
        self.forwarded = 0

    def forward(self, input_data: int):
        self.forwarded += 1
        return input_data + 1

    def num_forwarded(self):
        return self.forwarded

# Create an actor via ``remote`` API not ``bind`` API to avoid
# killing actors when a DAG is finished.
worker = Worker.remote()

with InputNode() as input_data:
    dag = MultiOutputNode([worker.forward.bind(input_data)])

# Actors are reused. The DAG definition doesn't include
# actor creation.
assert ray.get(dag.execute(1)) == [2]
assert ray.get(dag.execute(2)) == [3]
assert ray.get(dag.execute(3)) == [4]

# You can still use other actor methods via `remote` API.
assert ray.get(worker.num_forwarded.remote()) == 3

More resources#

For more application patterns and examples, see the following resources from other Ray libraries built on top of the Ray DAG API: