[asv] Add benchmark for ray.wait

python/benchmarks/benchmark_wait.py

@@ -0,0 +1,39 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+
+import ray
+
+
+def setup(*args):
+    if not hasattr(setup, "is_initialized"):
+        ray.init(num_workers=4, num_cpus=4)
+        setup.is_initialized = True
+
+
+@ray.remote
+def sleep(x):
+    time.sleep(x)
+
+
+class WaitSuite(object):
+    timeout = 10
+    timer = time.time
+
+    def time_wait_task(self):
+        ray.wait([sleep.remote(0.1)])
+
+    def time_wait_many_tasks(self, num_returns):
+        tasks = [sleep.remote(i / 5) for i in range(4)]
+        ray.wait(tasks, num_returns=num_returns)
+
+    time_wait_many_tasks.params = list(range(1, 4))
+    time_wait_many_tasks.param_names = ["num_returns"]
+
+    def time_wait_timeout(self, timeout):
+        ray.wait([sleep.remote(0.5)], timeout=timeout)
+
+    time_wait_timeout.params = [200, 800]
+    time_wait_timeout.param_names = ["timeout_ms"]

python/benchmarks/benchmarks.py

@@ -3,7 +3,6 @@
 from __future__ import print_function
 
 import ray
-import ray.test.test_functions as test_functions
 
 
 def setup():
@@ -12,6 +11,11 @@ def setup():
         setup.is_initialized = True
 
 
+@ray.remote
+def trivial_function():
+    return 1
+
+
 class TimeSuite(object):
     """An example benchmark."""
 
@@ -37,10 +41,10 @@ def mem_list(self):
 
 class MicroBenchmarkSuite(object):
     def time_submit(self):
-        test_functions.empty_function.remote()
+        trivial_function.remote()
 
     def time_submit_and_get(self):
-        x = test_functions.trivial_function.remote()
+        x = trivial_function.remote()
         ray.get(x)
 
     def time_put(self):