diff --git a/python/benchmarks/benchmark_task.py b/python/benchmarks/benchmark_task.py new file mode 100644 index 000000000..30a4bb8cb --- /dev/null +++ b/python/benchmarks/benchmark_task.py @@ -0,0 +1,61 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ray + + +def setup(): + if not hasattr(setup, "is_initialized"): + ray.init(num_workers=10, num_cpus=10, resources={"foo": 1}) + setup.is_initialized = True + + +def square(x): + return x * x + + +class TaskSuite(object): + timeout = 10 + + def setup(self): + self.square = ray.remote(square) + + def run_many_tasks(self): + ray.get([self.square.remote(i) for i in range(100)]) + + def run_task_dependency(self): + first_oid = self.square.remote(2) + second_oid = self.square.remote(first_oid) + ray.get(second_oid) + + def time_submit_task(self): + self.square.remote(2) + + def time_task_lifecycle(self): + ray.get(self.square.remote(2)) + + def peakmem_task_lifecycle(self): + ray.get(self.square.remote(2)) + + def time_run_many_tasks(self): + self.run_many_tasks() + + def peakmem_run_many_tasks(self): + self.run_many_tasks() + + def time_task_dependency(self): + self.run_task_dependency() + + def peakmem_task_dependency(self): + self.run_task_dependency() + + +class CPUTaskSuite(TaskSuite): + def setup(self): + self.square = ray.remote(num_cpus=1)(square) + + +class CustomResourceTaskSuite(TaskSuite): + def setup(self): + self.square = ray.remote(resources={"foo": 1})(square)