337 lines
11 KiB
Python
337 lines
11 KiB
Python
"""This is the script for `ray microbenchmark`."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import multiprocessing
|
|
|
|
import ray
|
|
import ray.experimental.channel as ray_channel
|
|
from ray._common.utils import (
|
|
get_or_create_event_loop,
|
|
)
|
|
from ray._private.ray_microbenchmark_helpers import asyncio_timeit, timeit
|
|
from ray._private.test_utils import get_actor_node_id
|
|
from ray.dag import InputNode, MultiOutputNode
|
|
from ray.dag.compiled_dag_node import CompiledDAG
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@ray.remote
|
|
class DAGActor:
|
|
def echo(self, x):
|
|
return x
|
|
|
|
def echo_multiple(self, *x):
|
|
return x
|
|
|
|
|
|
def check_optimized_build():
|
|
if not ray._raylet.OPTIMIZED:
|
|
msg = (
|
|
"WARNING: Unoptimized build! "
|
|
"To benchmark an optimized build, try:\n"
|
|
"\tbazel run -c opt //:gen_ray_pkg\n"
|
|
"You can also make this permanent by adding\n"
|
|
"\tbuild --compilation_mode=opt\n"
|
|
"to your user-wide ~/.bazelrc file. "
|
|
"(Do not add this to the project-level .bazelrc file.)"
|
|
)
|
|
logger.warning(msg)
|
|
|
|
|
|
def create_driver_actor():
|
|
return CompiledDAG.DAGDriverProxyActor.options(
|
|
label_selector={
|
|
ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
|
|
}
|
|
).remote()
|
|
|
|
|
|
def main(results=None):
|
|
results = results or []
|
|
loop = get_or_create_event_loop()
|
|
|
|
check_optimized_build()
|
|
|
|
print("Tip: set TESTS_TO_RUN='pattern' to run a subset of benchmarks")
|
|
|
|
#################################################
|
|
# Perf tests for channels, used in compiled DAGs.
|
|
#################################################
|
|
ray.init()
|
|
|
|
def put_channel_small(chans, do_get=False):
|
|
for chan in chans:
|
|
chan.write(b"0")
|
|
if do_get:
|
|
chan.read()
|
|
|
|
@ray.remote
|
|
class ChannelReader:
|
|
def ready(self):
|
|
return
|
|
|
|
def read(self, chans):
|
|
while True:
|
|
for chan in chans:
|
|
chan.read()
|
|
|
|
driver_actor = create_driver_actor()
|
|
driver_node = get_actor_node_id(driver_actor)
|
|
chans = [ray_channel.Channel(None, [(driver_actor, driver_node)], 1000)]
|
|
results += timeit(
|
|
"[unstable] local put:local get, single channel calls",
|
|
lambda: put_channel_small(chans, do_get=True),
|
|
)
|
|
|
|
reader = ChannelReader.remote()
|
|
reader_node = get_actor_node_id(reader)
|
|
chans = [ray_channel.Channel(None, [(reader, reader_node)], 1000)]
|
|
ray.get(reader.ready.remote())
|
|
reader.read.remote(chans)
|
|
results += timeit(
|
|
"[unstable] local put:1 remote get, single channel calls",
|
|
lambda: put_channel_small(chans),
|
|
)
|
|
ray.kill(reader)
|
|
|
|
n_cpu = multiprocessing.cpu_count() // 2
|
|
print(f"Testing multiple readers/channels, n={n_cpu}")
|
|
|
|
reader_and_node_list = []
|
|
for _ in range(n_cpu):
|
|
reader = ChannelReader.remote()
|
|
reader_node = get_actor_node_id(reader)
|
|
reader_and_node_list.append((reader, reader_node))
|
|
chans = [ray_channel.Channel(None, reader_and_node_list, 1000)]
|
|
ray.get([reader.ready.remote() for reader, _ in reader_and_node_list])
|
|
for reader, _ in reader_and_node_list:
|
|
reader.read.remote(chans)
|
|
results += timeit(
|
|
"[unstable] local put:n remote get, single channel calls",
|
|
lambda: put_channel_small(chans),
|
|
)
|
|
for reader, _ in reader_and_node_list:
|
|
ray.kill(reader)
|
|
|
|
reader = ChannelReader.remote()
|
|
reader_node = get_actor_node_id(reader)
|
|
chans = [
|
|
ray_channel.Channel(None, [(reader, reader_node)], 1000) for _ in range(n_cpu)
|
|
]
|
|
ray.get(reader.ready.remote())
|
|
reader.read.remote(chans)
|
|
results += timeit(
|
|
"[unstable] local put:1 remote get, n channels calls",
|
|
lambda: put_channel_small(chans),
|
|
)
|
|
ray.kill(reader)
|
|
|
|
reader_and_node_list = []
|
|
for _ in range(n_cpu):
|
|
reader = ChannelReader.remote()
|
|
reader_node = get_actor_node_id(reader)
|
|
reader_and_node_list.append((reader, reader_node))
|
|
chans = [
|
|
ray_channel.Channel(None, [reader_and_node_list[i]], 1000) for i in range(n_cpu)
|
|
]
|
|
ray.get([reader.ready.remote() for reader, _ in reader_and_node_list])
|
|
for chan, reader_node_tuple in zip(chans, reader_and_node_list):
|
|
reader = reader_node_tuple[0]
|
|
reader.read.remote([chan])
|
|
results += timeit(
|
|
"[unstable] local put:n remote get, n channels calls",
|
|
lambda: put_channel_small(chans),
|
|
)
|
|
for reader, _ in reader_and_node_list:
|
|
ray.kill(reader)
|
|
|
|
# Tests for compiled DAGs.
|
|
|
|
def _exec(dag, num_args=1, payload_size=1):
|
|
output_ref = dag.execute(*[b"x" * payload_size for _ in range(num_args)])
|
|
ray.get(output_ref)
|
|
|
|
async def exec_async(tag):
|
|
async def _exec_async():
|
|
fut = await compiled_dag.execute_async(b"x")
|
|
if not isinstance(fut, list):
|
|
await fut
|
|
else:
|
|
await asyncio.gather(*fut)
|
|
|
|
return await asyncio_timeit(
|
|
tag,
|
|
_exec_async,
|
|
)
|
|
|
|
# Single-actor DAG calls
|
|
|
|
a = DAGActor.remote()
|
|
with InputNode() as inp:
|
|
dag = a.echo.bind(inp)
|
|
|
|
results += timeit(
|
|
"[unstable] single-actor DAG calls", lambda: ray.get(dag.execute(b"x"))
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
"[unstable] compiled single-actor DAG calls", lambda: _exec(compiled_dag)
|
|
)
|
|
del a
|
|
|
|
# Single-actor asyncio DAG calls
|
|
|
|
a = DAGActor.remote()
|
|
with InputNode() as inp:
|
|
dag = a.echo.bind(inp)
|
|
compiled_dag = dag.experimental_compile(enable_asyncio=True)
|
|
results += loop.run_until_complete(
|
|
exec_async(
|
|
"[unstable] compiled single-actor asyncio DAG calls",
|
|
)
|
|
)
|
|
del a
|
|
|
|
# Scatter-gather DAG calls
|
|
|
|
n_cpu = multiprocessing.cpu_count() // 2
|
|
actors = [DAGActor.remote() for _ in range(n_cpu)]
|
|
with InputNode() as inp:
|
|
dag = MultiOutputNode([a.echo.bind(inp) for a in actors])
|
|
results += timeit(
|
|
f"[unstable] scatter-gather DAG calls, n={n_cpu} actors",
|
|
lambda: ray.get(dag.execute(b"x")),
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
f"[unstable] compiled scatter-gather DAG calls, n={n_cpu} actors",
|
|
lambda: _exec(compiled_dag),
|
|
)
|
|
|
|
# Scatter-gather asyncio DAG calls
|
|
|
|
actors = [DAGActor.remote() for _ in range(n_cpu)]
|
|
with InputNode() as inp:
|
|
dag = MultiOutputNode([a.echo.bind(inp) for a in actors])
|
|
compiled_dag = dag.experimental_compile(enable_asyncio=True)
|
|
results += loop.run_until_complete(
|
|
exec_async(
|
|
f"[unstable] compiled scatter-gather asyncio DAG calls, n={n_cpu} actors",
|
|
)
|
|
)
|
|
|
|
# Chain DAG calls
|
|
|
|
actors = [DAGActor.remote() for _ in range(n_cpu)]
|
|
with InputNode() as inp:
|
|
dag = inp
|
|
for a in actors:
|
|
dag = a.echo.bind(dag)
|
|
results += timeit(
|
|
f"[unstable] chain DAG calls, n={n_cpu} actors",
|
|
lambda: ray.get(dag.execute(b"x")),
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
f"[unstable] compiled chain DAG calls, n={n_cpu} actors",
|
|
lambda: _exec(compiled_dag),
|
|
)
|
|
|
|
# Chain asyncio DAG calls
|
|
|
|
actors = [DAGActor.remote() for _ in range(n_cpu)]
|
|
with InputNode() as inp:
|
|
dag = inp
|
|
for a in actors:
|
|
dag = a.echo.bind(dag)
|
|
compiled_dag = dag.experimental_compile(enable_asyncio=True)
|
|
results += loop.run_until_complete(
|
|
exec_async(f"[unstable] compiled chain asyncio DAG calls, n={n_cpu} actors")
|
|
)
|
|
|
|
# Multiple args with small payloads
|
|
|
|
n_actors = 8
|
|
assert (
|
|
n_cpu > n_actors
|
|
), f"n_cpu ({n_cpu}) must be greater than n_actors ({n_actors})"
|
|
|
|
actors = [DAGActor.remote() for _ in range(n_actors)]
|
|
with InputNode() as inp:
|
|
dag = MultiOutputNode([actors[i].echo.bind(inp[i]) for i in range(n_actors)])
|
|
payload_size = 1
|
|
results += timeit(
|
|
f"[unstable] multiple args with small payloads DAG calls, n={n_actors} actors",
|
|
lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_actors)])),
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
f"[unstable] compiled multiple args with small payloads DAG calls, "
|
|
f"n={n_actors} actors",
|
|
lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
|
|
)
|
|
|
|
# Multiple args with medium payloads
|
|
|
|
actors = [DAGActor.remote() for _ in range(n_actors)]
|
|
with InputNode() as inp:
|
|
dag = MultiOutputNode([actors[i].echo.bind(inp[i]) for i in range(n_actors)])
|
|
payload_size = 1024 * 1024
|
|
results += timeit(
|
|
f"[unstable] multiple args with medium payloads DAG calls, n={n_actors} actors",
|
|
lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_actors)])),
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
"[unstable] compiled multiple args with medium payloads DAG calls, "
|
|
f"n={n_actors} actors",
|
|
lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
|
|
)
|
|
|
|
# Multiple args with large payloads
|
|
|
|
actors = [DAGActor.remote() for _ in range(n_actors)]
|
|
with InputNode() as inp:
|
|
dag = MultiOutputNode([actors[i].echo.bind(inp[i]) for i in range(n_actors)])
|
|
payload_size = 10 * 1024 * 1024
|
|
results += timeit(
|
|
f"[unstable] multiple args with large payloads DAG calls, n={n_actors} actors",
|
|
lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_actors)])),
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
"[unstable] compiled multiple args with large payloads DAG calls, "
|
|
f"n={n_actors} actors",
|
|
lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
|
|
)
|
|
|
|
# Worst case for multiple arguments: a single actor takes all the arguments
|
|
# with small payloads.
|
|
|
|
actor = DAGActor.remote()
|
|
n_args = 8
|
|
with InputNode() as inp:
|
|
dag = actor.echo_multiple.bind(*[inp[i] for i in range(n_args)])
|
|
payload_size = 1
|
|
results += timeit(
|
|
"[unstable] single-actor with all args with small payloads DAG calls, "
|
|
"n=1 actors",
|
|
lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_args)])),
|
|
)
|
|
compiled_dag = dag.experimental_compile()
|
|
results += timeit(
|
|
"[unstable] single-actor with all args with small payloads DAG calls, "
|
|
"n=1 actors",
|
|
lambda: _exec(compiled_dag, num_args=n_args, payload_size=payload_size),
|
|
)
|
|
|
|
ray.shutdown()
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|