Files
2026-07-13 13:17:40 +08:00

139 lines
4.2 KiB
Python

import time
import click
import tqdm
from many_nodes_tests.dashboard_test import DashboardTestAtScale
import ray
import ray._common.test_utils
import ray._private.test_utils as test_utils
from ray._private.state_api_test_utils import (
StateAPICallSpec,
periodic_invoke_state_apis_with_actor,
summarize_worker_startup_time,
)
from ray.util.state import summarize_tasks
sleep_time = 300
def test_max_running_tasks(num_tasks):
cpus_per_task = 0.25
@ray.remote(num_cpus=cpus_per_task)
def task():
time.sleep(sleep_time)
def time_up(start_time):
return time.time() - start_time >= sleep_time
refs = [task.remote() for _ in tqdm.trange(num_tasks, desc="Launching tasks")]
max_cpus = ray.cluster_resources()["CPU"]
min_cpus_available = max_cpus
start_time = time.time()
for _ in tqdm.trange(int(sleep_time / 0.1), desc="Waiting"):
try:
cur_cpus = ray.available_resources().get("CPU", 0)
min_cpus_available = min(min_cpus_available, cur_cpus)
except Exception:
# There are race conditions `.get` can fail if a new heartbeat
# comes at the same time.
pass
if time_up(start_time):
print(f"Time up for sleeping {sleep_time} seconds")
break
time.sleep(0.1)
# There are some relevant magic numbers in this check. 10k tasks each
# require 1/4 cpus. Therefore, ideally 2.5k cpus will be used.
used_cpus = max_cpus - min_cpus_available
err_str = f"Only {used_cpus}/{max_cpus} cpus used."
# 1500 tasks. Note that it is a pretty low threshold, and the
# performance should be tracked via perf dashboard.
threshold = num_tasks * cpus_per_task * 0.60
print(f"{used_cpus}/{max_cpus} used.")
assert used_cpus > threshold, err_str
for _ in tqdm.trange(num_tasks, desc="Ensuring all tasks have finished"):
done, refs = ray.wait(refs)
assert ray.get(done[0]) is None
return used_cpus
def no_resource_leaks():
return test_utils.no_resource_leaks_excluding_node_resources()
@click.command()
@click.option("--num-tasks", required=True, type=int, help="Number of tasks to launch.")
def test(num_tasks):
addr = ray.init(address="auto")
ray._common.test_utils.wait_for_condition(no_resource_leaks)
monitor_actor = test_utils.monitor_memory_usage()
dashboard_test = DashboardTestAtScale(addr)
def not_none(res):
return res is not None
api_caller = periodic_invoke_state_apis_with_actor(
apis=[StateAPICallSpec(summarize_tasks, not_none)],
call_interval_s=4,
print_result=True,
)
start_time = time.time()
used_cpus = test_max_running_tasks(num_tasks)
end_time = time.time()
ray.get(monitor_actor.stop_run.remote())
used_gb, usage = ray.get(monitor_actor.get_peak_memory_info.remote())
print(f"Peak memory usage: {round(used_gb, 2)}GB")
print(f"Peak memory usage per processes:\n {usage}")
ray.get(api_caller.stop.remote())
del api_caller
del monitor_actor
ray._common.test_utils.wait_for_condition(no_resource_leaks)
try:
summarize_worker_startup_time()
except Exception as e:
print("Failed to summarize worker startup time.")
print(e)
rate = num_tasks / (end_time - start_time - sleep_time)
print(
f"Success! Started {num_tasks} tasks in {end_time - start_time}s. "
f"({rate} tasks/s)"
)
results = {
"tasks_per_second": rate,
"num_tasks": num_tasks,
"time": end_time - start_time,
"used_cpus": used_cpus,
"_peak_memory": round(used_gb, 2),
"_peak_process_memory": usage,
"perf_metrics": [
{
"perf_metric_name": "tasks_per_second",
"perf_metric_value": rate,
"perf_metric_type": "THROUGHPUT",
},
{
"perf_metric_name": "used_cpus_by_deadline",
"perf_metric_value": used_cpus,
"perf_metric_type": "THROUGHPUT",
},
],
}
dashboard_test.update_release_test_result(results)
test_utils.safe_write_to_results_json(results)
if __name__ == "__main__":
test()