wehub-resource-sync
135 lines
3.7 KiB
Python
135 lines
3.7 KiB
Python
import argparse
|
|
from math import floor
|
|
from time import sleep, time
|
|
|
|
import ray
|
|
import ray._private.test_utils as test_utils
|
|
from ray._private.test_utils import safe_write_to_results_json
|
|
|
|
|
|
@ray.remote
|
|
def simple_task(t):
|
|
sleep(t)
|
|
|
|
|
|
@ray.remote
|
|
class SimpleActor:
|
|
def __init__(self, job=None):
|
|
self._job = job
|
|
|
|
def ready(self):
|
|
return
|
|
|
|
def do_job(self):
|
|
if self._job is not None:
|
|
self._job()
|
|
|
|
|
|
def start_tasks(num_task, num_cpu_per_task, task_duration):
|
|
ray.get(
|
|
[
|
|
simple_task.options(num_cpus=num_cpu_per_task).remote(task_duration)
|
|
for _ in range(num_task)
|
|
]
|
|
)
|
|
|
|
|
|
def measure(f):
|
|
start = time()
|
|
ret = f()
|
|
end = time()
|
|
return (end - start, ret)
|
|
|
|
|
|
def start_actor(num_actors, num_actors_per_nodes, job):
|
|
resources = {"node": floor(1.0 / num_actors_per_nodes)}
|
|
submission_cost, actors = measure(
|
|
lambda: [
|
|
SimpleActor.options(resources=resources, num_cpus=0).remote(job)
|
|
for _ in range(num_actors)
|
|
]
|
|
)
|
|
ready_cost, _ = measure(lambda: ray.get([actor.ready.remote() for actor in actors]))
|
|
actor_job_cost, _ = measure(
|
|
lambda: ray.get([actor.do_job.remote() for actor in actors])
|
|
)
|
|
return (submission_cost, ready_cost, actor_job_cost)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(prog="Test Scheduling")
|
|
# Task workloads
|
|
parser.add_argument(
|
|
"--total-num-task", type=int, help="Total number of tasks.", required=False
|
|
)
|
|
parser.add_argument(
|
|
"--num-cpu-per-task",
|
|
type=int,
|
|
help="Resources needed for tasks.",
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
"--task-duration-s",
|
|
type=int,
|
|
help="How long does each task execute.",
|
|
required=False,
|
|
default=1,
|
|
)
|
|
|
|
# Actor workloads
|
|
parser.add_argument(
|
|
"--total-num-actors", type=int, help="Total number of actors.", required=True
|
|
)
|
|
parser.add_argument(
|
|
"--num-actors-per-nodes",
|
|
type=int,
|
|
help="How many actors to allocate for each nodes.",
|
|
required=True,
|
|
)
|
|
|
|
ray.init(address="auto")
|
|
monitor_actor = test_utils.monitor_memory_usage()
|
|
|
|
total_cpus_per_node = [node["Resources"].get("CPU", 0) for node in ray.nodes()]
|
|
num_nodes = len(total_cpus_per_node)
|
|
total_cpus = sum(total_cpus_per_node)
|
|
|
|
args = parser.parse_args()
|
|
job = None
|
|
if args.total_num_task is not None:
|
|
if args.num_cpu_per_task is None:
|
|
args.num_cpu_per_task = floor(1.0 * total_cpus / args.total_num_task)
|
|
job = lambda: start_tasks( # noqa: E731
|
|
args.total_num_task, args.num_cpu_per_task, args.task_duration_s
|
|
)
|
|
|
|
submission_cost, ready_cost, actor_job_cost = start_actor(
|
|
args.total_num_actors, args.num_actors_per_nodes, job
|
|
)
|
|
|
|
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}")
|
|
del monitor_actor
|
|
|
|
result = {
|
|
"total_num_task": args.total_num_task,
|
|
"num_cpu_per_task": args.num_cpu_per_task,
|
|
"task_duration_s": args.task_duration_s,
|
|
"total_num_actors": args.total_num_actors,
|
|
"num_actors_per_nodes": args.num_actors_per_nodes,
|
|
"num_nodes": num_nodes,
|
|
"total_cpus": total_cpus,
|
|
"submission_cost": submission_cost,
|
|
"ready_cost": ready_cost,
|
|
"actor_job_cost": actor_job_cost,
|
|
"_peak_memory": round(used_gb, 2),
|
|
"_peak_process_memory": usage,
|
|
"_runtime": submission_cost + ready_cost + actor_job_cost,
|
|
}
|
|
|
|
safe_write_to_results_json(result)
|
|
|
|
print(result)
|