[xray] Implements ray.wait

src/local_scheduler/lib/python/local_scheduler_extension.cc

@@ -183,12 +183,15 @@ static PyObject *PyLocalSchedulerClient_wait(PyObject *self, PyObject *args) {
   PyObject *py_object_ids;
   int num_returns;
   int64_t timeout_ms;
-  int wait_local;
+  PyObject *py_wait_local;
 
-  if (!PyArg_ParseTuple(args, "Oili", &py_object_ids, &num_returns, &timeout_ms,
-                        &wait_local)) {
+  if (!PyArg_ParseTuple(args, "OilO", &py_object_ids, &num_returns, &timeout_ms,
+                        &py_wait_local)) {
     return NULL;
   }
+
+  bool wait_local = PyObject_IsTrue(py_wait_local);
+
   // Convert object ids.
   PyObject *iter = PyObject_GetIter(py_object_ids);
   if (!iter) {

src/ray/object_manager/object_directory.cc

@@ -120,9 +120,8 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i
   listener_state.callbacks.emplace(callback_id, callback);
   // Immediately notify of found object locations.
   if (!listener_state.current_object_locations.empty()) {
-    std::vector<ClientID> client_id_vec(
-        listener_state.current_object_locations.begin(),
-        listener_state.current_object_locations.end());
+    std::vector<ClientID> client_id_vec(listener_state.current_object_locations.begin(),
+                                        listener_state.current_object_locations.end());
     callback(client_id_vec, object_id);
   }
   return status;

src/ray/object_manager/object_manager.cc

@@ -377,7 +377,7 @@ ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
   wait_state.object_id_order = object_ids;
   wait_state.timeout_ms = timeout_ms;
   wait_state.num_required_objects = num_required_objects;
-  for (auto &object_id : object_ids) {
+  for (const auto &object_id : object_ids) {
     if (local_objects_.count(object_id) > 0) {
       wait_state.found.insert(object_id);
     } else {
@@ -393,7 +393,7 @@ ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
     // we obtain information about all given objects, regardless of their location.
     // This is required to ensure we do not bias returning locally available objects
     // as ready whenever Wait is invoked with a mixture of local and remote objects.
-    for (auto &object_id : wait_state.remaining) {
+    for (const auto &object_id : wait_state.remaining) {
       // Lookup remaining objects.
       wait_state.requested_objects.insert(object_id);
       RAY_CHECK_OK(object_directory_->LookupLocations(
@@ -422,7 +422,7 @@ void ObjectManager::AllWaitLookupsComplete(const UniqueID &wait_id) {
     WaitComplete(wait_id);
   } else {
     // Subscribe to objects in order to ensure Wait-related tests are deterministic.
-    for (auto &object_id : wait_state.object_id_order) {
+    for (const auto &object_id : wait_state.object_id_order) {
       if (wait_state.remaining.count(object_id) == 0) {
         continue;
       }
@@ -470,7 +470,7 @@ void ObjectManager::WaitComplete(const UniqueID &wait_id) {
     RAY_CHECK(wait_state.timeout_ms > 0 || wait_state.timeout_ms == -1);
   }
   // Unsubscribe to any objects that weren't found in the time allotted.
-  for (auto &object_id : wait_state.requested_objects) {
+  for (const auto &object_id : wait_state.requested_objects) {
     RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations(wait_id, object_id));
   }
   // Cancel the timer. This is okay even if the timer hasn't been started.
@@ -480,7 +480,7 @@ void ObjectManager::WaitComplete(const UniqueID &wait_id) {
   // Order objects according to input order.
   std::vector<ObjectID> found;
   std::vector<ObjectID> remaining;
-  for (auto item : wait_state.object_id_order) {
+  for (const auto item : wait_state.object_id_order) {
     if (found.size() < wait_state.num_required_objects &&
         wait_state.found.count(item) > 0) {
       found.push_back(item);

src/ray/object_manager/test/object_manager_test.cc

@@ -285,12 +285,12 @@ class TestObjectManagerCommands : public TestObjectManager {
     ObjectID object_2 = WriteDataToClient(client2, data_size);
     UniqueID sub_id = ray::ObjectID::from_random();
 
-    get_object_directory(*server1).SubscribeObjectLocations(
+    RAY_CHECK_OK(get_object_directory(*server1).SubscribeObjectLocations(
         sub_id, object_1,
         [this, sub_id, object_1, object_2](const std::vector<ray::ClientID> &,
                                            const ray::ObjectID &object_id) {
           TestWaitCallbacks2(sub_id, object_1, object_2);
-        });
+        }));
   }
 
   void TestWaitCallbacks2(UniqueID sub_id, ObjectID object_1, ObjectID object_2) {
@@ -313,7 +313,8 @@ class TestObjectManagerCommands : public TestObjectManager {
           RAY_CHECK(found.size() == 1);
           // There's nothing more to test. A check will fail if unexpected behavior is
           // triggered.
-          get_object_directory(*server1).UnsubscribeObjectLocations(sub_id, object_1);
+          RAY_CHECK_OK(get_object_directory(*server1).UnsubscribeObjectLocations(
+              sub_id, object_1));
           NextWaitTest();
         }));
   }