[xray] Object manager retries Pull requests (#2630)

* Move all ObjectManager members to bottom of class def

* Better Pull requests
- suppress duplicate Pulls
- retry the Pull at the next client after a timeout
- cancel a Pull if the object no longer appears on any clients

* increase object manager Pull timeout

* Make the component failure test harder.

* note

* Notify SubscribeObjectLocations caller of empty list

* Address melih's comments

* Fix wait...

* Make component failure test easier for legacy ray

* lint

src/common/state/ray_config.h

@@ -149,10 +149,7 @@ class RayConfig {
         max_tasks_to_spillback_(10),
         actor_creation_num_spillbacks_warning_(100),
         node_manager_forward_task_retry_timeout_milliseconds_(1000),
-        // TODO: Setting this to large values results in latency, which needs to
-        // be addressed. This timeout is often on the critical path for object
-        // transfers.
-        object_manager_pull_timeout_ms_(20),
+        object_manager_pull_timeout_ms_(100),
         object_manager_push_timeout_ms_(10000),
         object_manager_default_chunk_size_(1000000),
         num_workers_per_process_(1) {}

src/ray/object_manager/object_directory.cc

@@ -53,14 +53,15 @@ void ObjectDirectory::RegisterBackend() {
     std::vector<ClientID> client_id_vec =
         UpdateObjectLocations(object_id_listener_pair->second.current_object_locations,
                               location_history, gcs_client_->client_table());
-    if (!client_id_vec.empty()) {
-      // Copy the callbacks so that the callbacks can unsubscribe without interrupting
-      // looping over the callbacks.
-      auto callbacks = object_id_listener_pair->second.callbacks;
-      // Call all callbacks associated with the object id locations we have received.
-      for (const auto &callback_pair : callbacks) {
-        callback_pair.second(client_id_vec, object_id);
-      }
+    // Copy the callbacks so that the callbacks can unsubscribe without interrupting
+    // looping over the callbacks.
+    auto callbacks = object_id_listener_pair->second.callbacks;
+    // Call all callbacks associated with the object id locations we have
+    // received.  This notifies the client even if the list of locations is
+    // empty, since this may indicate that the objects have been evicted from
+    // all nodes.
+    for (const auto &callback_pair : callbacks) {
+      callback_pair.second(client_id_vec, object_id);
     }
   };
   RAY_CHECK_OK(gcs_client_->object_table().Subscribe(
@@ -131,12 +132,12 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i
     return ray::Status::OK();
   }
   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());
-    callback(client_id_vec, object_id);
-  }
+  // Immediately notify of object locations. This notifies the client even if
+  // the list of locations is empty, since this may indicate that the objects
+  // have been evicted from all nodes.
+  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_directory.h

@@ -58,11 +58,13 @@ class ObjectDirectoryInterface {
                                       const OnLocationsFound &callback) = 0;
 
   /// Subscribe to be notified of locations (ClientID) of the given object.
-  /// The callback will be invoked whenever locations are obtained for the
-  /// specified object. The callback provided to this method may fire immediately,
-  /// within the call to this method, if any other listener is subscribed to the same
-  /// object: This occurs when location data for the object has already been obtained.
-  ///
+  /// The callback will be invoked with the complete list of known locations
+  /// whenever the set of locations changes. The callback will also be fired if
+  /// the list of known locations is empty. The callback provided to this
+  /// method may fire immediately, within the call to this method, if any other
+  /// listener is subscribed to the same object: This occurs when location data
+  /// for the object has already been obtained.
+  //
   /// \param callback_id The id associated with the specified callback. This is
   /// needed when UnsubscribeObjectLocations is called.
   /// \param object_id The required object's ObjectID.

src/ray/object_manager/object_manager.cc

@@ -34,10 +34,8 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   RAY_CHECK(config_.max_sends > 0);
   RAY_CHECK(config_.max_receives > 0);
   main_service_ = &main_service;
-  store_notification_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
-    NotifyDirectoryObjectAdd(object_info);
-    HandleUnfulfilledPushRequests(object_info);
-  });
+  store_notification_.SubscribeObjAdded(
+      [this](const ObjectInfoT &object_info) { HandleObjectAdded(object_info); });
   store_notification_.SubscribeObjDeleted(
       [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); });
   StartIOService();
@@ -60,10 +58,8 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   RAY_CHECK(config_.max_receives > 0);
   // TODO(hme) Client ID is never set with this constructor.
   main_service_ = &main_service;
-  store_notification_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
-    NotifyDirectoryObjectAdd(object_info);
-    HandleUnfulfilledPushRequests(object_info);
-  });
+  store_notification_.SubscribeObjAdded(
+      [this](const ObjectInfoT &object_info) { HandleObjectAdded(object_info); });
   store_notification_.SubscribeObjDeleted(
       [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); });
   StartIOService();
@@ -97,15 +93,13 @@ void ObjectManager::StopIOService() {
   }
 }
 
-void ObjectManager::NotifyDirectoryObjectAdd(const ObjectInfoT &object_info) {
+void ObjectManager::HandleObjectAdded(const ObjectInfoT &object_info) {
+  // Notify the object directory that the object has been added to this node.
   ObjectID object_id = ObjectID::from_binary(object_info.object_id);
   local_objects_[object_id] = object_info;
   ray::Status status =
       object_directory_->ReportObjectAdded(object_id, client_id_, object_info);
-}
 
-void ObjectManager::HandleUnfulfilledPushRequests(const ObjectInfoT &object_info) {
-  ObjectID object_id = ObjectID::from_binary(object_info.object_id);
   // Handle the unfulfilled_push_requests_ which contains the push request that is not
   // completed due to unsatisfied local objects.
   auto iter = unfulfilled_push_requests_.find(object_id);
@@ -120,6 +114,10 @@ void ObjectManager::HandleUnfulfilledPushRequests(const ObjectInfoT &object_info
     }
     unfulfilled_push_requests_.erase(iter);
   }
+
+  // The object is local, so we no longer need to Pull it from a remote
+  // manager. Cancel any outstanding Pull requests for this object.
+  CancelPull(object_id);
 }
 
 void ObjectManager::NotifyDirectoryObjectDeleted(const ObjectID &object_id) {
@@ -145,38 +143,107 @@ ray::Status ObjectManager::Pull(const ObjectID &object_id) {
     RAY_LOG(ERROR) << object_id << " attempted to pull an object that's already local.";
     return ray::Status::OK();
   }
-  ray::Status status_code = object_directory_->SubscribeObjectLocations(
+  if (pull_requests_.find(object_id) != pull_requests_.end()) {
+    return ray::Status::OK();
+  }
+
+  pull_requests_.emplace(object_id, PullRequest());
+  // Subscribe to object notifications. A notification will be received every
+  // time the set of client IDs for the object changes. Notifications will also
+  // be received if the list of locations is empty. The set of client IDs has
+  // no ordering guarantee between notifications.
+  return object_directory_->SubscribeObjectLocations(
       object_directory_pull_callback_id_, object_id,
       [this](const std::vector<ClientID> &client_ids, const ObjectID &object_id) {
-        RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations(
-            object_directory_pull_callback_id_, object_id));
-        GetLocationsSuccess(client_ids, object_id);
+        // Exit if the Pull request has already been fulfilled or canceled.
+        auto it = pull_requests_.find(object_id);
+        if (it == pull_requests_.end()) {
+          return;
+        }
+        // Reset the list of clients that are now expected to have the object.
+        // NOTE(swang): Since we are overwriting the previous list of clients,
+        // we may end up sending a duplicate request to the same client as
+        // before.
+        it->second.client_locations = client_ids;
+        if (it->second.client_locations.empty()) {
+          // The object locations are now empty, so we should wait for the next
+          // notification about a new object location.  Cancel the timer until
+          // the next Pull attempt since there are no more clients to try.
+          if (it->second.retry_timer != nullptr) {
+            it->second.retry_timer->cancel();
+            it->second.timer_set = false;
+          }
+        } else {
+          // New object locations were found.
+          if (!it->second.timer_set) {
+            // The timer was not set, which means that we weren't trying any
+            // clients. We now have some clients to try, so begin trying to
+            // Pull from one.  If we fail to receive an object within the pull
+            // timeout, then this will try the rest of the clients in the list
+            // in succession.
+            TryPull(object_id);
+          }
+        }
       });
-  return status_code;
 }
 
-void ObjectManager::GetLocationsSuccess(const std::vector<ray::ClientID> &client_ids,
-                                        const ray::ObjectID &object_id) {
-  if (local_objects_.count(object_id) == 0) {
-    // Only pull objects that aren't local.
-    RAY_CHECK(!client_ids.empty());
-    ClientID client_id = client_ids.front();
-    Pull(object_id, client_id);
-  }
-}
-
-void ObjectManager::Pull(const ObjectID &object_id, const ClientID &client_id) {
-  // Check if object is already local.
-  if (local_objects_.count(object_id) != 0) {
-    RAY_LOG(ERROR) << object_id << " attempted to pull an object that's already local.";
+void ObjectManager::TryPull(const ObjectID &object_id) {
+  auto it = pull_requests_.find(object_id);
+  if (it == pull_requests_.end()) {
     return;
   }
-  // Check if we're pulling from self.
+
+  // The timer should never fire if there are no expected client locations.
+  RAY_CHECK(!it->second.client_locations.empty());
+  RAY_CHECK(local_objects_.count(object_id) == 0);
+
+  // Get the next client to try.
+  const ClientID client_id = std::move(it->second.client_locations.back());
+  it->second.client_locations.pop_back();
   if (client_id == client_id_) {
+    // If we're trying to pull from ourselves, skip this client and try the
+    // next one.
     RAY_LOG(ERROR) << client_id_ << " attempted to pull an object from itself.";
-    return;
+    const ClientID client_id = std::move(it->second.client_locations.back());
+    it->second.client_locations.pop_back();
+    RAY_CHECK(client_id != client_id_);
   }
+
+  // Try pulling from the client.
   PullEstablishConnection(object_id, client_id);
+
+  // If there are more clients to try, try them in succession, with a timeout
+  // in between each try.
+  if (!it->second.client_locations.empty()) {
+    if (it->second.retry_timer == nullptr) {
+      // Set the timer if we haven't already.
+      it->second.retry_timer = std::unique_ptr<boost::asio::deadline_timer>(
+          new boost::asio::deadline_timer(*main_service_));
+    }
+
+    // Wait for a timeout. If we receive the object or a caller Cancels the
+    // Pull within the timeout, then nothing will happen. Otherwise, the timer
+    // will fire and the next client in the list will be tried.
+    boost::posix_time::milliseconds retry_timeout(config_.pull_timeout_ms);
+    it->second.retry_timer->expires_from_now(retry_timeout);
+    it->second.retry_timer->async_wait(
+        [this, object_id](const boost::system::error_code &error) {
+          if (!error) {
+            // Try the Pull from the next client.
+            TryPull(object_id);
+          } else {
+            // Check that the error was due to the timer being canceled.
+            RAY_CHECK(error == boost::asio::error::operation_aborted);
+          }
+        });
+    // Record that we set the timer until the next attempt.
+    it->second.timer_set = true;
+  } else {
+    // The timer is not reset since there are no more clients to try. Go back
+    // to waiting for more notifications. Once we receive a new object location
+    // from the object directory, then the Pull will be retried.
+    it->second.timer_set = false;
+  }
 };
 
 void ObjectManager::PullEstablishConnection(const ObjectID &object_id,
@@ -370,10 +437,15 @@ ray::Status ObjectManager::SendObjectData(const ObjectID &object_id,
   return status;
 }
 
-ray::Status ObjectManager::Cancel(const ObjectID &object_id) {
-  ray::Status status = object_directory_->UnsubscribeObjectLocations(
-      object_directory_pull_callback_id_, object_id);
-  return status;
+void ObjectManager::CancelPull(const ObjectID &object_id) {
+  auto it = pull_requests_.find(object_id);
+  if (it == pull_requests_.end()) {
+    return;
+  }
+
+  RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations(
+      object_directory_pull_callback_id_, object_id));
+  pull_requests_.erase(it);
 }
 
 ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
@@ -481,22 +553,26 @@ void ObjectManager::SubscribeRemainingWaitObjects(const UniqueID &wait_id) {
       RAY_CHECK_OK(object_directory_->SubscribeObjectLocations(
           wait_id, object_id, [this, wait_id](const std::vector<ClientID> &client_ids,
                                               const ObjectID &subscribe_object_id) {
-            auto object_id_wait_state = active_wait_requests_.find(wait_id);
-            // We never expect to handle a subscription notification for a wait that has
-            // already completed.
-            RAY_CHECK(object_id_wait_state != active_wait_requests_.end());
-            auto &wait_state = object_id_wait_state->second;
-            RAY_CHECK(wait_state.remaining.erase(subscribe_object_id));
-            wait_state.found.insert(subscribe_object_id);
-            wait_state.requested_objects.erase(subscribe_object_id);
-            RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations(
-                wait_id, subscribe_object_id));
-            if (wait_state.found.size() >= wait_state.num_required_objects) {
-              WaitComplete(wait_id);
+            if (!client_ids.empty()) {
+              auto object_id_wait_state = active_wait_requests_.find(wait_id);
+              // We never expect to handle a subscription notification for a wait that has
+              // already completed.
+              RAY_CHECK(object_id_wait_state != active_wait_requests_.end());
+              auto &wait_state = object_id_wait_state->second;
+              RAY_CHECK(wait_state.remaining.erase(subscribe_object_id));
+              wait_state.found.insert(subscribe_object_id);
+              wait_state.requested_objects.erase(subscribe_object_id);
+              RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations(
+                  wait_id, subscribe_object_id));
+              if (wait_state.found.size() >= wait_state.num_required_objects) {
+                WaitComplete(wait_id);
+              }
             }
           }));
     }
     if (wait_state.timeout_ms != -1) {
+      auto timeout = boost::posix_time::milliseconds(wait_state.timeout_ms);
+      wait_state.timeout_timer->expires_from_now(timeout);
       wait_state.timeout_timer->async_wait(
           [this, wait_id](const boost::system::error_code &error_code) {
             if (error_code.value() != 0) {

src/ray/object_manager/object_manager.h

@@ -50,7 +50,7 @@ struct ObjectManagerConfig {
 class ObjectManagerInterface {
  public:
   virtual ray::Status Pull(const ObjectID &object_id) = 0;
-  virtual ray::Status Cancel(const ObjectID &object_id) = 0;
+  virtual void CancelPull(const ObjectID &object_id) = 0;
   virtual ~ObjectManagerInterface(){};
 };
 
@@ -104,20 +104,22 @@ class ObjectManager : public ObjectManagerInterface {
   /// \return Void.
   void Push(const ObjectID &object_id, const ClientID &client_id);
 
-  /// Pull an object from ClientID. Returns UniqueID asociated with
-  /// an invocation of this method.
+  /// Pull an object from ClientID.
   ///
   /// \param object_id The object's object id.
   /// \return Status of whether the pull request successfully initiated.
   ray::Status Pull(const ObjectID &object_id);
 
-  /// Discover ClientID via ObjectDirectory, then pull object
-  /// from ClientID associated with ObjectID.
+  /// Try to Pull an object from one of its expected client locations. If there
+  /// are more client locations to try after this attempt, then this method
+  /// will try each of the other clients in succession, with a timeout between
+  /// each attempt. If the object is received or if the Pull is Canceled before
+  /// the timeout, then no more Pull requests for this object will be sent
+  /// to other node managers until TryPull is called again.
   ///
   /// \param object_id The object's object id.
-  /// \param client_id The remote node's client id.
   /// \return Void.
-  void Pull(const ObjectID &object_id, const ClientID &client_id);
+  void TryPull(const ObjectID &object_id);
 
   /// Add a connection to a remote object manager.
   /// This is invoked by an external server.
@@ -136,11 +138,12 @@ class ObjectManager : public ObjectManagerInterface {
   void ProcessClientMessage(std::shared_ptr<TcpClientConnection> &conn,
                             int64_t message_type, const uint8_t *message);
 
-  /// Cancels all requests (Push/Pull) associated with the given ObjectID.
+  /// Cancels all requests (Push/Pull) associated with the given ObjectID. This
+  /// method is idempotent.
   ///
   /// \param object_id The ObjectID.
-  /// \return Status of whether requests were successfully cancelled.
-  ray::Status Cancel(const ObjectID &object_id);
+  /// \return Void.
+  void CancelPull(const ObjectID &object_id);
 
   /// Callback definition for wait.
   using WaitCallback = std::function<void(const std::vector<ray::ObjectID> &found,
@@ -163,45 +166,12 @@ class ObjectManager : public ObjectManagerInterface {
  private:
   friend class TestObjectManager;
 
-  ClientID client_id_;
-  const ObjectManagerConfig config_;
-  std::unique_ptr<ObjectDirectoryInterface> object_directory_;
-  ObjectStoreNotificationManager store_notification_;
-  ObjectBufferPool buffer_pool_;
-
-  /// This runs on a thread pool dedicated to sending objects.
-  boost::asio::io_service send_service_;
-  /// This runs on a thread pool dedicated to receiving objects.
-  boost::asio::io_service receive_service_;
-
-  /// Weak reference to main service. We ensure this object is destroyed before
-  /// main_service_ is stopped.
-  boost::asio::io_service *main_service_;
-
-  /// Used to create "work" for send_service_.
-  /// Without this, if send_service_ has no more sends to process, it will stop.
-  boost::asio::io_service::work send_work_;
-  /// Used to create "work" for receive_service_.
-  /// Without this, if receive_service_ has no more receives to process, it will stop.
-  boost::asio::io_service::work receive_work_;
-
-  /// Runs the send service, which handle
-  /// all outgoing object transfers.
-  std::vector<std::thread> send_threads_;
-  /// Runs the receive service, which handle
-  /// all incoming object transfers.
-  std::vector<std::thread> receive_threads_;
-
-  /// Connection pool for reusing outgoing connections to remote object managers.
-  ConnectionPool connection_pool_;
-
-  /// Cache of locally available objects.
-  std::unordered_map<ObjectID, ObjectInfoT> local_objects_;
-
-  /// This is used as the callback identifier in Pull for
-  /// SubscribeObjectLocations. We only need one identifier because we never need to
-  /// subscribe multiple times to the same object during Pull.
-  UniqueID object_directory_pull_callback_id_ = UniqueID::from_random();
+  struct PullRequest {
+    PullRequest() : retry_timer(nullptr), timer_set(false), client_locations() {}
+    std::unique_ptr<boost::asio::deadline_timer> retry_timer;
+    bool timer_set;
+    std::vector<ClientID> client_locations;
+  };
 
   struct WaitState {
     WaitState(asio::io_service &service, int64_t timeout_ms, const WaitCallback &callback)
@@ -228,9 +198,6 @@ class ObjectManager : public ObjectManagerInterface {
     uint64_t num_required_objects;
   };
 
-  /// A set of active wait requests.
-  std::unordered_map<UniqueID, WaitState> active_wait_requests_;
-
   /// Creates a wait request and adds it to active_wait_requests_.
   ray::Status AddWaitRequest(const UniqueID &wait_id,
                              const std::vector<ObjectID> &object_ids, int64_t timeout_ms,
@@ -247,39 +214,25 @@ class ObjectManager : public ObjectManagerInterface {
   /// Completion handler for Wait.
   void WaitComplete(const UniqueID &wait_id);
 
-  /// Maintains a map of push requests that have not been fulfilled due to an object not
-  /// being local. Objects are removed from this map after push_timeout_ms have elapsed.
-  std::unordered_map<
-      ObjectID,
-      std::unordered_map<ClientID, std::unique_ptr<boost::asio::deadline_timer>>>
-      unfulfilled_push_requests_;
-
   /// Handle starting, running, and stopping asio io_service.
   void StartIOService();
   void RunSendService();
   void RunReceiveService();
   void StopIOService();
 
-  /// Register object add with directory.
-  void NotifyDirectoryObjectAdd(const ObjectInfoT &object_info);
+  /// Handle an object being added to this node. This adds the object to the
+  /// directory, pushes the object to other nodes if necessary, and cancels any
+  /// outstanding Pull requests for the object.
+  void HandleObjectAdded(const ObjectInfoT &object_info);
 
   /// Register object remove with directory.
   void NotifyDirectoryObjectDeleted(const ObjectID &object_id);
 
-  /// Handle any push requests that were made before an object was available.
-  /// This is invoked when an "object added" notification is received from the store.
-  void HandleUnfulfilledPushRequests(const ObjectInfoT &object_info);
-
   /// Part of an asynchronous sequence of Pull methods.
   /// Uses an existing connection or creates a connection to ClientID.
   /// Executes on main_service_ thread.
   void PullEstablishConnection(const ObjectID &object_id, const ClientID &client_id);
 
-  /// Private callback implementation for success on get location. Called from
-  /// ObjectDirectory.
-  void GetLocationsSuccess(const std::vector<ray::ClientID> &client_ids,
-                           const ray::ObjectID &object_id);
-
   /// Synchronously send a pull request via remote object manager connection.
   /// Executes on main_service_ thread.
   ray::Status PullSendRequest(const ObjectID &object_id,
@@ -326,6 +279,58 @@ class ObjectManager : public ObjectManagerInterface {
                         const uint8_t *message);
   /// Handle Push task timeout.
   void HandlePushTaskTimeout(const ObjectID &object_id, const ClientID &client_id);
+
+  ClientID client_id_;
+  const ObjectManagerConfig config_;
+  std::unique_ptr<ObjectDirectoryInterface> object_directory_;
+  ObjectStoreNotificationManager store_notification_;
+  ObjectBufferPool buffer_pool_;
+
+  /// This runs on a thread pool dedicated to sending objects.
+  boost::asio::io_service send_service_;
+  /// This runs on a thread pool dedicated to receiving objects.
+  boost::asio::io_service receive_service_;
+
+  /// Weak reference to main service. We ensure this object is destroyed before
+  /// main_service_ is stopped.
+  boost::asio::io_service *main_service_;
+
+  /// Used to create "work" for send_service_.
+  /// Without this, if send_service_ has no more sends to process, it will stop.
+  boost::asio::io_service::work send_work_;
+  /// Used to create "work" for receive_service_.
+  /// Without this, if receive_service_ has no more receives to process, it will stop.
+  boost::asio::io_service::work receive_work_;
+
+  /// Runs the send service, which handle
+  /// all outgoing object transfers.
+  std::vector<std::thread> send_threads_;
+  /// Runs the receive service, which handle
+  /// all incoming object transfers.
+  std::vector<std::thread> receive_threads_;
+
+  /// Connection pool for reusing outgoing connections to remote object managers.
+  ConnectionPool connection_pool_;
+
+  /// Cache of locally available objects.
+  std::unordered_map<ObjectID, ObjectInfoT> local_objects_;
+
+  /// This is used as the callback identifier in Pull for
+  /// SubscribeObjectLocations. We only need one identifier because we never need to
+  /// subscribe multiple times to the same object during Pull.
+  UniqueID object_directory_pull_callback_id_ = UniqueID::from_random();
+
+  /// A set of active wait requests.
+  std::unordered_map<UniqueID, WaitState> active_wait_requests_;
+
+  /// Maintains a map of push requests that have not been fulfilled due to an object not
+  /// being local. Objects are removed from this map after push_timeout_ms have elapsed.
+  std::unordered_map<
+      ObjectID,
+      std::unordered_map<ClientID, std::unique_ptr<boost::asio::deadline_timer>>>
+      unfulfilled_push_requests_;
+
+  std::unordered_map<ObjectID, PullRequest> pull_requests_;
 };
 
 }  // namespace ray

src/ray/object_manager/test/object_manager_test.cc

@@ -284,9 +284,11 @@ class TestObjectManager : public TestObjectManagerBase {
 
     RAY_CHECK_OK(server1->object_manager_.object_directory_->SubscribeObjectLocations(
         sub_id, object_1,
-        [this, sub_id, object_1, object_2](const std::vector<ray::ClientID> &,
+        [this, sub_id, object_1, object_2](const std::vector<ray::ClientID> &clients,
                                            const ray::ObjectID &object_id) {
-          TestWaitWhileSubscribed(sub_id, object_1, object_2);
+          if (!clients.empty()) {
+            TestWaitWhileSubscribed(sub_id, object_1, object_2);
+          }
         }));
   }
 

src/ray/raylet/task_dependency_manager.cc

@@ -65,7 +65,7 @@ void TaskDependencyManager::HandleRemoteDependencyCanceled(const ObjectID &objec
   if (!required) {
     auto it = required_objects_.find(object_id);
     if (it != required_objects_.end()) {
-      RAY_CHECK_OK(object_manager_.Cancel(object_id));
+      object_manager_.CancelPull(object_id);
       reconstruction_policy_.Cancel(object_id);
       required_objects_.erase(it);
     }

src/ray/raylet/task_dependency_manager_test.cc

@@ -16,7 +16,7 @@ using ::testing::_;
 class MockObjectManager : public ObjectManagerInterface {
  public:
   MOCK_METHOD1(Pull, ray::Status(const ObjectID &object_id));
-  MOCK_METHOD1(Cancel, ray::Status(const ObjectID &object_id));
+  MOCK_METHOD1(CancelPull, void(const ObjectID &object_id));
 };
 
 class MockReconstructionPolicy : public ReconstructionPolicyInterface {
@@ -119,7 +119,7 @@ TEST_F(TaskDependencyManagerTest, TestSimpleTask) {
 
   // All arguments should be canceled as they become available locally.
   for (const auto &argument_id : arguments) {
-    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+    EXPECT_CALL(object_manager_mock_, CancelPull(argument_id));
     EXPECT_CALL(reconstruction_policy_mock_, Cancel(argument_id));
   }
   // For each argument except the last, tell the task dependency manager that
@@ -156,7 +156,7 @@ TEST_F(TaskDependencyManagerTest, TestDuplicateSubscribe) {
 
   // All arguments should be canceled as they become available locally.
   for (const auto &argument_id : arguments) {
-    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+    EXPECT_CALL(object_manager_mock_, CancelPull(argument_id));
     EXPECT_CALL(reconstruction_policy_mock_, Cancel(argument_id));
   }
   // For each argument except the last, tell the task dependency manager that
@@ -191,7 +191,7 @@ TEST_F(TaskDependencyManagerTest, TestMultipleTasks) {
   }
 
   // Tell the task dependency manager that the object is local.
-  EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+  EXPECT_CALL(object_manager_mock_, CancelPull(argument_id));
   EXPECT_CALL(reconstruction_policy_mock_, Cancel(argument_id));
   auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(argument_id);
   // Check that all tasks are now ready to run.
@@ -213,7 +213,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskChain) {
   // locally queued task.
   EXPECT_CALL(object_manager_mock_, Pull(_)).Times(0);
   EXPECT_CALL(reconstruction_policy_mock_, ListenAndMaybeReconstruct(_)).Times(0);
-  EXPECT_CALL(object_manager_mock_, Cancel(_)).Times(0);
+  EXPECT_CALL(object_manager_mock_, CancelPull(_)).Times(0);
   EXPECT_CALL(reconstruction_policy_mock_, Cancel(_)).Times(0);
   for (const auto &task : tasks) {
     // Subscribe to each of the tasks' arguments.
@@ -279,7 +279,7 @@ TEST_F(TaskDependencyManagerTest, TestDependentPut) {
 
   // The put object should be considered local as soon as the task that creates
   // it is pending execution.
-  EXPECT_CALL(object_manager_mock_, Cancel(put_id));
+  EXPECT_CALL(object_manager_mock_, CancelPull(put_id));
   EXPECT_CALL(reconstruction_policy_mock_, Cancel(put_id));
   EXPECT_CALL(gcs_mock_, Add(_, task1.GetTaskSpecification().TaskId(), _, _));
   task_dependency_manager_.TaskPending(task1);
@@ -312,7 +312,7 @@ TEST_F(TaskDependencyManagerTest, TestTaskForwarding) {
 
   // Simulate the task executing on a remote node and its return value
   // appearing locally.
-  EXPECT_CALL(object_manager_mock_, Cancel(return_id));
+  EXPECT_CALL(object_manager_mock_, CancelPull(return_id));
   EXPECT_CALL(reconstruction_policy_mock_, Cancel(return_id));
   auto ready_tasks = task_dependency_manager_.HandleObjectLocal(return_id);
   // Check that the task that we kept is now ready to run.
@@ -341,7 +341,7 @@ TEST_F(TaskDependencyManagerTest, TestEviction) {
   // Tell the task dependency manager that each of the arguments is now
   // available.
   for (const auto &argument_id : arguments) {
-    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+    EXPECT_CALL(object_manager_mock_, CancelPull(argument_id));
     EXPECT_CALL(reconstruction_policy_mock_, Cancel(argument_id));
   }
   for (size_t i = 0; i < arguments.size(); i++) {
@@ -379,7 +379,7 @@ TEST_F(TaskDependencyManagerTest, TestEviction) {
   // Tell the task dependency manager that each of the arguments is available
   // again.
   for (const auto &argument_id : arguments) {
-    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+    EXPECT_CALL(object_manager_mock_, CancelPull(argument_id));
     EXPECT_CALL(reconstruction_policy_mock_, Cancel(argument_id));
   }
   for (size_t i = 0; i < arguments.size(); i++) {

test/component_failures_test.py

@@ -138,6 +138,8 @@ class ComponentFailureTest(unittest.TestCase):
 
     def _testComponentFailed(self, component_type):
         """Kill a component on all worker nodes and check workload succeeds."""
+        # Raylet is able to pass a harder failure test than legacy ray.
+        use_raylet = os.environ.get("RAY_USE_XRAY") == "1"
 
         # Start with 4 workers and 4 cores.
         num_local_schedulers = 4
@@ -149,41 +151,80 @@ class ComponentFailureTest(unittest.TestCase):
             num_cpus=[num_workers_per_scheduler] * num_local_schedulers,
             redirect_output=True)
 
-        # Submit many tasks with many dependencies.
-        @ray.remote
-        def f(x):
-            return x
+        if use_raylet:
+            # Submit many tasks with many dependencies.
+            @ray.remote
+            def f(x):
+                return x
 
-        x = 1
-        for _ in range(1000):
-            x = f.remote(x)
-        ray.get(x)
+            @ray.remote
+            def g(*xs):
+                return 1
 
-        @ray.remote
-        def g(*xs):
-            return 1
+            # Kill the component on all nodes except the head node as the tasks
+            # execute. Do this in a loop while submitting tasks between each
+            # component failure.
+            # NOTE(swang): Legacy ray hangs on this test if the plasma manager
+            # is killed.
+            time.sleep(0.1)
+            components = ray.services.all_processes[component_type]
+            for process in components[1:]:
+                # Submit a round of tasks with many dependencies.
+                x = 1
+                for _ in range(1000):
+                    x = f.remote(x)
 
-        xs = [g.remote(1)]
-        for _ in range(100):
-            xs.append(g.remote(*xs))
-            xs.append(g.remote(1))
+                xs = [g.remote(1)]
+                for _ in range(100):
+                    xs.append(g.remote(*xs))
+                    xs.append(g.remote(1))
 
-        # Kill the component on all nodes except the head node as the tasks
-        # execute.
-        time.sleep(0.1)
-        components = ray.services.all_processes[component_type]
-        for process in components[1:]:
-            process.terminate()
-            time.sleep(1)
+                # Kill a component on one of the nodes.
+                process.terminate()
+                time.sleep(1)
+                process.kill()
+                process.wait()
+                assert not process.poll() is None
 
-        for process in components[1:]:
-            process.kill()
-            process.wait()
-            assert not process.poll() is None
+                # Make sure that we can still get the objects after the
+                # executing tasks died.
+                ray.get(x)
+                ray.get(xs)
+        else:
 
-        # Make sure that we can still get the objects after the executing tasks
-        # died.
-        ray.get(xs)
+            @ray.remote
+            def f(x, j):
+                time.sleep(0.2)
+                return x
+
+            # Submit more tasks than there are workers so that all workers and
+            # cores are utilized.
+            object_ids = [
+                f.remote(i, 0) for i in range(num_workers_per_scheduler *
+                                              num_local_schedulers)
+            ]
+            object_ids += [f.remote(object_id, 1) for object_id in object_ids]
+            object_ids += [f.remote(object_id, 2) for object_id in object_ids]
+
+            # Kill the component on all nodes except the head node as the tasks
+            # execute.
+            time.sleep(0.1)
+            components = ray.services.all_processes[component_type]
+            for process in components[1:]:
+                process.terminate()
+                time.sleep(1)
+
+            for process in components[1:]:
+                process.kill()
+                process.wait()
+                assert not process.poll() is None
+
+            # Make sure that we can still get the objects after the executing
+            # tasks died.
+            results = ray.get(object_ids)
+            expected_results = 4 * list(
+                range(num_workers_per_scheduler * num_local_schedulers))
+            assert results == expected_results
 
     def check_components_alive(self, component_type, check_component_alive):
         """Check that a given component type is alive on all worker nodes.