wehub-resource-sync
112 lines
2.7 KiB
Python
112 lines
2.7 KiB
Python
# This workload tests many drivers using the same cluster.
|
|
import time
|
|
import argparse
|
|
|
|
import ray
|
|
from ray._common.test_utils import run_string_as_driver
|
|
from ray._private.test_utils import safe_write_to_results_json
|
|
|
|
|
|
def update_progress(result):
|
|
result["last_update"] = time.time()
|
|
safe_write_to_results_json(result)
|
|
|
|
|
|
ray.init()
|
|
|
|
nodes = [node["NodeID"] for node in ray.nodes()]
|
|
assert len(nodes) == 4
|
|
|
|
# Run the workload.
|
|
|
|
# Define a driver script that runs a few tasks and actors on each node in the
|
|
# cluster.
|
|
driver_script = f"""
|
|
import ray
|
|
|
|
ray.init()
|
|
|
|
nodes = "{','.join(nodes)}".split(',')
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def f():
|
|
return 1
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor(object):
|
|
def method(self):
|
|
return 1
|
|
|
|
for _ in range(5):
|
|
for node in nodes:
|
|
assert ray.get(
|
|
f.options(label_selector={{ray._raylet.RAY_NODE_ID_KEY: node}}).remote()) == 1
|
|
actor = Actor.options(label_selector={{ray._raylet.RAY_NODE_ID_KEY: node}}).remote()
|
|
assert ray.get(actor.method.remote()) == 1
|
|
|
|
print("success")
|
|
"""
|
|
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def run_driver():
|
|
output = run_string_as_driver(driver_script, encode="utf-8")
|
|
assert "success" in output
|
|
|
|
|
|
iteration = 0
|
|
running_ids = [
|
|
run_driver.options(label_selector={ray._raylet.RAY_NODE_ID_KEY: node}).remote()
|
|
for node in nodes
|
|
]
|
|
start_time = time.time()
|
|
previous_time = start_time
|
|
|
|
parser = argparse.ArgumentParser(prog="Many Drivers long running tests")
|
|
parser.add_argument(
|
|
"--iteration-num", type=int, help="How many iterations to run", required=False
|
|
)
|
|
parser.add_argument(
|
|
"--smoke-test",
|
|
action="store_true",
|
|
help="Whether or not the test is smoke test.",
|
|
default=False,
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
iteration_num = args.iteration_num
|
|
if args.smoke_test:
|
|
iteration_num = 400
|
|
while True:
|
|
if iteration_num is not None and iteration_num < iteration:
|
|
break
|
|
# Wait for a driver to finish and start a new driver.
|
|
[ready_id], running_ids = ray.wait(running_ids, num_returns=1)
|
|
ray.get(ready_id)
|
|
|
|
running_ids.append(
|
|
run_driver.options(
|
|
label_selector={ray._raylet.RAY_NODE_ID_KEY: nodes[iteration % len(nodes)]}
|
|
).remote()
|
|
)
|
|
|
|
new_time = time.time()
|
|
print(
|
|
"Iteration {}:\n"
|
|
" - Iteration time: {}.\n"
|
|
" - Absolute time: {}.\n"
|
|
" - Total elapsed time: {}.".format(
|
|
iteration, new_time - previous_time, new_time, new_time - start_time
|
|
)
|
|
)
|
|
update_progress(
|
|
{
|
|
"iteration": iteration,
|
|
"iteration_time": new_time - previous_time,
|
|
"absolute_time": new_time,
|
|
"elapsed_time": new_time - start_time,
|
|
}
|
|
)
|
|
previous_time = new_time
|
|
iteration += 1
|