Implement select to enable waiting for a specific number of remote objects to be ready. (#369)

lib/python/ray/__init__.py

@@ -11,7 +11,7 @@ if hasattr(ctypes, "windll"):
 
 import config
 import serialization
-from worker import scheduler_info, visualize_computation_graph, task_info, init, connect, disconnect, get, put, remote, kill_workers, restart_workers_local
+from worker import scheduler_info, visualize_computation_graph, task_info, init, connect, disconnect, get, put, select, remote, kill_workers, restart_workers_local
 from worker import Reusable, reusables
 from libraylib import SCRIPT_MODE, WORKER_MODE, PYTHON_MODE, SILENT_MODE
 from libraylib import ObjectID

lib/python/ray/worker.py

@@ -819,6 +819,36 @@ def put(value, worker=global_worker):
   worker.put_object(objectid, value)
   return objectid
 
+def select(objectids, num_objects=0, worker=global_worker):
+  """Return a list of the indices of the objects that are ready.
+
+  If num_objects is 0, the function immediately returns the indices of all
+  objects that are ready. If it is set, the function waits until that number of
+  objects is ready and returns that exact number of objectids.
+
+  Args:
+    objectids (List[ray.ObjectID]): List of objectids for objects that may or
+      may not be ready.
+    num_objects (int): The number of indices that should be returned.
+
+  Returns:
+    List of indices in the original list of objects that are ready.
+  """
+  check_connected(worker)
+  if num_objects > len(objectids):
+    raise Exception("num_objects cannot be greater than len(objectids), num_objects is {}, and len(objectids) is {}.".format(num_objects, len(objectids)))
+  ready_ids = raylib.ray_select(worker.handle, objectids)
+  # Polls scheduler until enough objects are ready.
+  while len(ready_ids) < num_objects:
+    ready_ids = raylib.ray_select(worker.handle, objectids)
+    time.sleep(0.1)
+  if num_objects != 0:
+    # Return indices for exactly the requested number of objects.
+    return ready_ids[:num_objects]
+  else:
+    # Return indices for all objects that are ready.
+    return ready_ids
+
 def kill_workers(worker=global_worker):
   """Kill all of the workers in the cluster. This does not kill drivers.