diff --git a/src/common/state/ray_config.h b/src/common/state/ray_config.h index 155f65c84..d4ae3cea8 100644 --- a/src/common/state/ray_config.h +++ b/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) {} diff --git a/src/ray/object_manager/object_directory.cc b/src/ray/object_manager/object_directory.cc index 86d9440f5..80a3f990d 100644 --- a/src/ray/object_manager/object_directory.cc +++ b/src/ray/object_manager/object_directory.cc @@ -53,14 +53,15 @@ void ObjectDirectory::RegisterBackend() { std::vector 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 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 client_id_vec(listener_state.current_object_locations.begin(), + listener_state.current_object_locations.end()); + callback(client_id_vec, object_id); return status; } diff --git a/src/ray/object_manager/object_directory.h b/src/ray/object_manager/object_directory.h index 549f01230..0de499135 100644 --- a/src/ray/object_manager/object_directory.h +++ b/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. diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 809d0e0b7..a7c59bc1d 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/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 &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 &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( + 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 &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 &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) { diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index 14630b483..7ff2c4930 100644 --- a/src/ray/object_manager/object_manager.h +++ b/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 &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 &found, @@ -163,45 +166,12 @@ class ObjectManager : public ObjectManagerInterface { private: friend class TestObjectManager; - ClientID client_id_; - const ObjectManagerConfig config_; - std::unique_ptr 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 send_threads_; - /// Runs the receive service, which handle - /// all incoming object transfers. - std::vector receive_threads_; - - /// Connection pool for reusing outgoing connections to remote object managers. - ConnectionPool connection_pool_; - - /// Cache of locally available objects. - std::unordered_map 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 retry_timer; + bool timer_set; + std::vector 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 active_wait_requests_; - /// Creates a wait request and adds it to active_wait_requests_. ray::Status AddWaitRequest(const UniqueID &wait_id, const std::vector &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>> - 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 &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 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 send_threads_; + /// Runs the receive service, which handle + /// all incoming object transfers. + std::vector receive_threads_; + + /// Connection pool for reusing outgoing connections to remote object managers. + ConnectionPool connection_pool_; + + /// Cache of locally available objects. + std::unordered_map 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 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>> + unfulfilled_push_requests_; + + std::unordered_map pull_requests_; }; } // namespace ray diff --git a/src/ray/object_manager/test/object_manager_test.cc b/src/ray/object_manager/test/object_manager_test.cc index d3e165cad..52362a279 100644 --- a/src/ray/object_manager/test/object_manager_test.cc +++ b/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 &, + [this, sub_id, object_1, object_2](const std::vector &clients, const ray::ObjectID &object_id) { - TestWaitWhileSubscribed(sub_id, object_1, object_2); + if (!clients.empty()) { + TestWaitWhileSubscribed(sub_id, object_1, object_2); + } })); } diff --git a/src/ray/raylet/task_dependency_manager.cc b/src/ray/raylet/task_dependency_manager.cc index 0adc4e997..ddd00c95c 100644 --- a/src/ray/raylet/task_dependency_manager.cc +++ b/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); } diff --git a/src/ray/raylet/task_dependency_manager_test.cc b/src/ray/raylet/task_dependency_manager_test.cc index 32c339028..e62ef1dca 100644 --- a/src/ray/raylet/task_dependency_manager_test.cc +++ b/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++) { diff --git a/test/component_failures_test.py b/test/component_failures_test.py index 8948b8682..44f410834 100644 --- a/test/component_failures_test.py +++ b/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.