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

src/raylib.cc

@@ -897,6 +897,26 @@ static PyObject* request_object(PyObject* self, PyObject* args) {
   Py_RETURN_NONE;
 }
 
+static PyObject* ray_select(PyObject* self, PyObject* args) {
+  Worker* worker;
+  PyObject* objectids;
+  if (!PyArg_ParseTuple(args, "O&O", &PyObjectToWorker, &worker, &objectids)) {
+    return NULL;
+  }
+  std::vector<ObjectID> objectids_vec;
+  for (size_t i = 0; i < PyList_Size(objectids); ++i) {
+    ObjectID objectid;
+    PyObjectToObjectID(PyList_GetItem(objectids, i), &objectid);
+    objectids_vec.push_back(objectid);
+  }
+  std::vector<int> indices = worker->select(objectids_vec);
+  PyObject* result = PyList_New(indices.size());
+  for (size_t i = 0; i < indices.size(); ++i) {
+    PyList_SetItem(result, i, PyInt_FromLong(indices[i]));
+  }
+  return result;
+}
+
 static PyObject* alias_objectids(PyObject* self, PyObject* args) {
   Worker* worker;
   ObjectID alias_objectid;
@@ -1061,6 +1081,7 @@ static PyMethodDef RayLibMethods[] = {
  { "get_object", get_object, METH_VARARGS, "get protocol buffer object from the local object store" },
  { "get_objectid", get_objectid, METH_VARARGS, "register a new object reference with the scheduler" },
  { "request_object" , request_object, METH_VARARGS, "request an object to be delivered to the local object store" },
+ { "ray_select" , ray_select, METH_VARARGS, "checks the scheduler to see if a object can be gotten" },
  { "alias_objectids", alias_objectids, METH_VARARGS, "make two objectids refer to the same object" },
  { "wait_for_next_message", wait_for_next_message, METH_VARARGS, "get next message from scheduler (blocking)" },
  { "submit_task", submit_task, METH_VARARGS, "call a remote function" },

src/scheduler.cc

@@ -549,6 +549,21 @@ Status SchedulerService::ExportReusableVariable(ServerContext* context, const Ex
   return Status::OK;
 }
 
+Status SchedulerService::Select(ServerContext* context, const SelectRequest* request, SelectReply* reply) {
+  auto objtable = GET(objtable_);
+  for (int i = 0; i < request->objectids_size(); ++i) {
+    ObjectID objectid = request->objectids(i);
+    if (has_canonical_objectid(objectid)) {
+      ObjectID canonical_objectid = get_canonical_objectid(objectid);
+      RAY_CHECK_LT(canonical_objectid, objtable->size(), "Canonical_objectid is outside object table.");
+      if ((*objtable)[canonical_objectid].size() != 0) {
+        reply->add_indices(i);
+      }
+    }
+  }
+  return Status::OK;
+}
+
 void SchedulerService::deliver_object_async_if_necessary(ObjectID canonical_objectid, ObjStoreId from, ObjStoreId to) {
   bool object_present_or_in_transit;
   {

src/scheduler.h

@@ -78,6 +78,7 @@ public:
   Status ExportRemoteFunction(ServerContext* context, const ExportRemoteFunctionRequest* request, AckReply* reply) override;
   Status ExportReusableVariable(ServerContext* context, const ExportReusableVariableRequest* request, AckReply* reply) override;
   Status NotifyFailure(ServerContext*, const NotifyFailureRequest* request, AckReply* reply) override;
+  Status Select(ServerContext*, const SelectRequest* request, SelectReply* reply) override;
 
 #ifdef NDEBUG
   // If we've disabled assertions, then just use regular SynchronizedPtr to skip lock checking.

src/worker.cc

@@ -424,6 +424,22 @@ void Worker::task_info(ClientContext &context, TaskInfoRequest &request, TaskInf
   RAY_CHECK_GRPC(scheduler_stub_->TaskInfo(&context, request, &reply));
 }
 
+std::vector<int> Worker::select(std::vector<ObjectID>& objectids) {
+  RAY_CHECK(connected_, "Attempted to test if object was ready but failed.");
+  ClientContext context;
+  SelectRequest request;
+  SelectReply reply;
+  for (int i = 0; i < objectids.size(); ++i) {
+    request.add_objectids(objectids[i]);
+  }
+  RAY_CHECK_GRPC(scheduler_stub_->Select(&context, request, &reply));
+  std::vector<int> result;
+  for (int i = 0; i < reply.indices_size(); ++i) {
+    result.push_back(reply.indices(i));
+  }
+  return result;
+}
+
 bool Worker::export_remote_function(const std::string& function_name, const std::string& function) {
   RAY_CHECK(connected_, "Attempted to export function but failed.");
   ClientContext context;

src/worker.h

@@ -102,6 +102,8 @@ class Worker {
   void scheduler_info(ClientContext &context, SchedulerInfoRequest &request, SchedulerInfoReply &reply);
   // get task statuses from scheduler
   void task_info(ClientContext &context, TaskInfoRequest &request, TaskInfoReply &reply);
+  // gets indices of available objects
+  std::vector<int> select(std::vector<ObjectID>& objectids);
   // export function to workers
   bool export_remote_function(const std::string& function_name, const std::string& function);
   // export reusable variable to workers