diff --git a/BUILD.bazel b/BUILD.bazel index b24538f14..c23326d3d 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -813,6 +813,18 @@ cc_test( ], ) +cc_test( + name = "pull_manager_test", + srcs = [ + "src/ray/object_manager/test/pull_manager_test.cc", + ], + copts = COPTS, + deps = [ + ":raylet_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "push_manager_test", srcs = [ diff --git a/src/ray/common/ray_config_def.h b/src/ray/common/ray_config_def.h index 4d1110e20..82faa3be9 100644 --- a/src/ray/common/ray_config_def.h +++ b/src/ray/common/ray_config_def.h @@ -177,6 +177,10 @@ RAY_CONFIG(int64_t, worker_register_timeout_seconds, 30) RAY_CONFIG(int64_t, redis_db_connect_retries, 50) RAY_CONFIG(int64_t, redis_db_connect_wait_milliseconds, 100) +/// Timeout, in milliseconds, to wait before retrying a failed pull in the +/// ObjectManager. +RAY_CONFIG(int, object_manager_timer_freq_ms, 100) + /// Timeout, in milliseconds, to wait before retrying a failed pull in the /// ObjectManager. RAY_CONFIG(int, object_manager_pull_timeout_ms, 10000) diff --git a/src/ray/object_manager/common.h b/src/ray/object_manager/common.h index eacfb4b13..c7c531ffc 100644 --- a/src/ray/object_manager/common.h +++ b/src/ray/object_manager/common.h @@ -1,8 +1,12 @@ #pragma once +#include #include +#include "ray/common/id.h" +#include "ray/object_manager/format/object_manager_generated.h" + namespace ray { /// A callback to asynchronously spill objects when space is needed. @@ -18,4 +22,8 @@ using SpillObjectsCallback = /// A callback to call when space has been released. using SpaceReleasedCallback = std::function; +/// A callback to call when a spilled object needs to be returned to the object store. +using RestoreSpilledObjectCallback = std::function)>; + } // namespace ray diff --git a/src/ray/object_manager/object_directory.h b/src/ray/object_manager/object_directory.h index 0396309bd..7133d1e94 100644 --- a/src/ray/object_manager/object_directory.h +++ b/src/ray/object_manager/object_directory.h @@ -40,6 +40,11 @@ struct RemoteConnectionInfo { uint16_t port; }; +/// Callback for object location notifications. +using OnLocationsFound = + std::function &, const std::string &)>; + class ObjectDirectoryInterface { public: virtual ~ObjectDirectoryInterface() {} @@ -58,12 +63,7 @@ class ObjectDirectoryInterface { /// \return A vector of information for all connected remote object managers. virtual std::vector LookupAllRemoteConnections() const = 0; - /// Callback for object location notifications. - using OnLocationsFound = - std::function &, const std::string &)>; - - /// Lookup object locations. Callback may be invoked with empty list of node ids. + /// Lookup object locations. Callback may be invoked with empty list of client ids. /// /// \param object_id The object's ObjectID. /// \param callback Invoked with (possibly empty) list of node ids and object_id. diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 90380e1b0..2eb641d04 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -57,25 +57,47 @@ ObjectManager::ObjectManager(asio::io_service &main_service, const NodeID &self_ RestoreSpilledObjectCallback restore_spilled_object, SpillObjectsCallback spill_objects_callback, std::function object_store_full_callback) - : self_node_id_(self_node_id), + : main_service_(&main_service), + self_node_id_(self_node_id), config_(config), object_directory_(std::move(object_directory)), object_store_internal_(config, spill_objects_callback, object_store_full_callback), buffer_pool_(config_.store_socket_name, config_.object_chunk_size), rpc_work_(rpc_service_), - gen_(std::chrono::high_resolution_clock::now().time_since_epoch().count()), object_manager_server_("ObjectManager", config_.object_manager_port, config_.rpc_service_threads_number), object_manager_service_(rpc_service_, *this), client_call_manager_(main_service, config_.rpc_service_threads_number), - restore_spilled_object_(restore_spilled_object) { + restore_spilled_object_(restore_spilled_object), + pull_retry_timer_(*main_service_, + boost::posix_time::milliseconds(config.timer_freq_ms)) { RAY_CHECK(config_.rpc_service_threads_number > 0); - main_service_ = &main_service; + + const auto &object_is_local = [this](const ObjectID &object_id) { + return local_objects_.count(object_id) != 0; + }; + const auto &send_pull_request = [this](const ObjectID &object_id, + const NodeID &client_id) { + SendPullRequest(object_id, client_id); + }; + const auto &get_time = []() { return absl::GetCurrentTimeNanos() / 1e9; }; + pull_manager_.reset(new PullManager(self_node_id_, object_is_local, send_pull_request, + restore_spilled_object_, get_time, + config.pull_timeout_ms)); push_manager_.reset(new PushManager(/* max_chunks_in_flight= */ std::max( static_cast(1L), static_cast(config_.max_bytes_in_flight / config_.object_chunk_size)))); + pull_retry_timer_.async_wait([this](const boost::system::error_code &e) { + RAY_CHECK(!e) << "The raylet's object manager has failed unexpectedly with error: " + << e + << ". Please file a bug report on here: " + "https://github.com/ray-project/ray/issues"; + + Tick(); + }); + if (plasma::plasma_store_runner) { store_notification_ = std::make_shared(main_service); plasma::plasma_store_runner->SetNotificationListener(store_notification_); @@ -178,171 +200,48 @@ ray::Status ObjectManager::SubscribeObjDeleted( ray::Status ObjectManager::Pull(const ObjectID &object_id, const rpc::Address &owner_address) { - RAY_LOG(DEBUG) << "Pull on " << self_node_id_ << " of object " << object_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."; - return ray::Status::OK(); - } - if (pull_requests_.find(object_id) != pull_requests_.end()) { - RAY_LOG(DEBUG) << object_id << " has inflight pull_requests, skipping."; - return ray::Status::OK(); + if (!pull_manager_->Pull(object_id, owner_address)) { + // If we don't need to pull, the object is either already local or this is a duplicate + // request. + return Status::OK(); } - pull_requests_.emplace(object_id, PullRequest()); + const auto &callback = [this](const ObjectID &object_id, + const std::unordered_set &client_ids, + const std::string &spilled_url) { + pull_manager_->OnLocationChange(object_id, client_ids, spilled_url); + }; + // Subscribe to object notifications. A notification will be received every // time the set of node IDs for the object changes. Notifications will also // be received if the list of locations is empty. The set of node IDs has // no ordering guarantee between notifications. - return object_directory_->SubscribeObjectLocations( - object_directory_pull_callback_id_, object_id, owner_address, - [this](const ObjectID &object_id, const std::unordered_set &node_ids, - const std::string &spilled_url) { - // 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 nodes that are now expected to have the object. - // NOTE(swang): Since we are overwriting the previous list of nodes, - // we may end up sending a duplicate request to the same node as - // before. - it->second.node_locations = std::vector(node_ids.begin(), node_ids.end()); - if (!spilled_url.empty()) { - // Try to restore the spilled object. - restore_spilled_object_(object_id, spilled_url, - [this, object_id](const ray::Status &status) { - // Fall back to fetching from another object manager. - if (!status.ok()) { - TryPull(object_id); - } - }); - } else if (it->second.node_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 nodes to try. - if (it->second.retry_timer != nullptr) { - it->second.retry_timer->cancel(); - it->second.timer_set = false; - } - } else { - // New object locations were found, so begin trying to pull from a - // node. This will be called every time a new node location - // appears. - TryPull(object_id); - } - }); + return object_directory_->SubscribeObjectLocations(object_directory_pull_callback_id_, + object_id, owner_address, callback); } -void ObjectManager::TryPull(const ObjectID &object_id) { - auto it = pull_requests_.find(object_id); - if (it == pull_requests_.end()) { - return; - } - - auto &node_vector = it->second.node_locations; - - // The timer should never fire if there are no expected node locations. - if (node_vector.empty()) { - return; - } - - RAY_CHECK(local_objects_.count(object_id) == 0); - // Make sure that there is at least one node which is not the local node. - // TODO(rkn): It may actually be possible for this check to fail. - if (node_vector.size() == 1 && node_vector[0] == self_node_id_) { - RAY_LOG(WARNING) << "The object manager with ID " << self_node_id_ - << " is trying to pull object " << object_id - << " but the object table suggests that this object manager " - << "already has the object. The object may have been evicted. It is " - << "most likely due to memory pressure, object pull has been " - << "requested before object location is updated."; - it->second.timer_set = false; - return; - } - - // Choose a random node to pull the object from. - // Generate a random index. - std::uniform_int_distribution distribution(0, node_vector.size() - 1); - int node_index = distribution(gen_); - NodeID node_id = node_vector[node_index]; - // If the object manager somehow ended up choosing itself, choose a different - // object manager. - if (node_id == self_node_id_) { - std::swap(node_vector[node_index], node_vector[node_vector.size() - 1]); - node_vector.pop_back(); - RAY_LOG(WARNING) - << "The object manager with ID " << self_node_id_ << " is trying to pull object " - << object_id << " but the object table suggests that this object manager " - << "already has the object. It is most likely due to memory pressure, object " - << "pull has been requested before object location is updated."; - node_id = node_vector[node_index % node_vector.size()]; - RAY_CHECK(node_id != self_node_id_); - } - - RAY_LOG(DEBUG) << "Sending pull request from " << self_node_id_ << " to " << node_id - << " of object " << object_id; - - auto rpc_client = GetRpcClient(node_id); +void ObjectManager::SendPullRequest(const ObjectID &object_id, const NodeID &client_id) { + auto rpc_client = GetRpcClient(client_id); if (rpc_client) { - // Try pulling from the node. - rpc_service_.post([this, object_id, node_id, rpc_client]() { - SendPullRequest(object_id, node_id, rpc_client); + // Try pulling from the client. + rpc_service_.post([this, object_id, client_id, rpc_client]() { + rpc::PullRequest pull_request; + pull_request.set_object_id(object_id.Binary()); + pull_request.set_node_id(self_node_id_.Binary()); + + rpc_client->Pull(pull_request, [object_id, client_id](const Status &status, + const rpc::PullReply &reply) { + if (!status.ok()) { + RAY_LOG(WARNING) << "Send pull " << object_id << " request to client " + << client_id << " failed due to" << status.message(); + } + }); }); } else { RAY_LOG(ERROR) << "Couldn't send pull request from " << self_node_id_ << " to " - << node_id << " of object " << object_id + << client_id << " of object " << object_id << " , setup rpc connection failed."; } - - // If there are more nodes to try, try them in succession, with a timeout - // in between each try. - if (!it->second.node_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 node 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 node. - 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 nodes 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::SendPullRequest( - const ObjectID &object_id, const NodeID &node_id, - std::shared_ptr rpc_client) { - rpc::PullRequest pull_request; - pull_request.set_object_id(object_id.Binary()); - pull_request.set_node_id(self_node_id_.Binary()); - - rpc_client->Pull(pull_request, [object_id, node_id](const Status &status, - const rpc::PullReply &reply) { - if (!status.ok()) { - RAY_LOG(WARNING) << "Send pull " << object_id << " request to node " << node_id - << " failed due to" << status.message(); - } - }); } void ObjectManager::HandlePushTaskTimeout(const ObjectID &object_id, @@ -528,14 +427,13 @@ void ObjectManager::SendObjectChunk(const UniqueID &push_id, const ObjectID &obj } void ObjectManager::CancelPull(const ObjectID &object_id) { - auto it = pull_requests_.find(object_id); - if (it == pull_requests_.end()) { + if (!pull_manager_->CancelPull(object_id)) { + // We weren't tracking a pull request for this object, so there is nothing to cancel. return; } RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations( object_directory_pull_callback_id_, object_id)); - pull_requests_.erase(it); } ray::Status ObjectManager::Wait( @@ -898,7 +796,7 @@ std::string ObjectManager::DebugString() const { result << "\n- num local objects: " << local_objects_.size(); result << "\n- num active wait requests: " << active_wait_requests_.size(); result << "\n- num unfulfilled push requests: " << unfulfilled_push_requests_.size(); - result << "\n- num pull requests: " << pull_requests_.size(); + result << "\n- num pull requests: " << pull_manager_->NumActiveRequests(); result << "\n- num buffered profile events: " << profile_events_.size(); result << "\n- num chunks received total: " << num_chunks_received_total_; result << "\n- num chunks received failed: " << num_chunks_received_failed_; @@ -913,7 +811,9 @@ void ObjectManager::RecordMetrics() const { stats::ObjectStoreAvailableMemory().Record(config_.object_store_memory - used_memory_); stats::ObjectStoreUsedMemory().Record(used_memory_); stats::ObjectStoreLocalObjects().Record(local_objects_.size()); - stats::ObjectManagerPullRequests().Record(pull_requests_.size()); + stats::ObjectManagerPullRequests().Record(pull_manager_->NumActiveRequests()); } +void ObjectManager::Tick() { pull_manager_->Tick(); } + } // namespace ray diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index d7d630c01..75dbb5bd4 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -32,12 +32,14 @@ #include "ray/common/id.h" #include "ray/common/ray_config.h" #include "ray/common/status.h" +#include "ray/object_manager/common.h" #include "ray/object_manager/format/object_manager_generated.h" #include "ray/object_manager/notification/object_store_notification_manager_ipc.h" #include "ray/object_manager/object_buffer_pool.h" #include "ray/object_manager/object_directory.h" #include "ray/object_manager/ownership_based_object_directory.h" #include "ray/object_manager/plasma/store_runner.h" +#include "ray/object_manager/pull_manager.h" #include "ray/object_manager/push_manager.h" #include "ray/rpc/object_manager/object_manager_client.h" #include "ray/rpc/object_manager/object_manager_server.h" @@ -49,6 +51,8 @@ struct ObjectManagerConfig { /// from other object managers. If this is 0, the object manager will choose /// its own port. int object_manager_port; + /// The object manager's global timer frequency. + unsigned int timer_freq_ms; /// The time in milliseconds to wait before retrying a pull /// that fails due to node id lookup. unsigned int pull_timeout_ms; @@ -78,7 +82,6 @@ struct LocalObjectInfo { /// Information from the object store about the object. object_manager::protocol::ObjectInfoT object_info; }; - class ObjectStoreRunner { public: ObjectStoreRunner(const ObjectManagerConfig &config, @@ -171,9 +174,8 @@ class ObjectManager : public ObjectManagerInterface, /// Send pull request /// /// \param object_id Object id - /// \param node_id Remote server node id - void SendPullRequest(const ObjectID &object_id, const NodeID &node_id, - std::shared_ptr rpc_client); + /// \param client_id Remote server client id + void SendPullRequest(const ObjectID &object_id, const NodeID &client_id); /// Get the rpc client according to the node ID /// @@ -235,17 +237,6 @@ class ObjectManager : public ObjectManagerInterface, /// \return Status of whether the pull request successfully initiated. ray::Status Pull(const ObjectID &object_id, const rpc::Address &owner_address) override; - /// Try to Pull an object from one of its expected node locations. If there - /// are more node locations to try after this attempt, then this method - /// will try each of the other nodes 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. - /// \return Void. - void TryPull(const ObjectID &object_id); - /// Cancels all requests (Push/Pull) associated with the given ObjectID. This /// method is idempotent. /// @@ -293,16 +284,11 @@ class ObjectManager : public ObjectManagerInterface, /// Record metrics. void RecordMetrics() const; + void Tick(); + private: friend class TestObjectManager; - struct PullRequest { - PullRequest() : retry_timer(nullptr), timer_set(false), node_locations() {} - std::unique_ptr retry_timer; - bool timer_set; - std::vector node_locations; - }; - struct WaitState { WaitState(boost::asio::io_service &service, int64_t timeout_ms, const WaitCallback &callback) @@ -406,6 +392,10 @@ class ObjectManager : public ObjectManagerInterface, /// Handle Push task timeout. void HandlePushTaskTimeout(const ObjectID &object_id, const NodeID &node_id); + /// Weak reference to main service. We ensure this object is destroyed before + /// main_service_ is stopped. + boost::asio::io_service *main_service_; + NodeID self_node_id_; const ObjectManagerConfig config_; std::shared_ptr object_directory_; @@ -416,10 +406,6 @@ class ObjectManager : public ObjectManagerInterface, std::shared_ptr store_notification_; ObjectBufferPool buffer_pool_; - /// Weak reference to main service. We ensure this object is destroyed before - /// main_service_ is stopped. - boost::asio::io_service *main_service_; - /// Multi-thread asio service, deal with all outgoing and incoming RPC request. boost::asio::io_service rpc_service_; @@ -448,10 +434,6 @@ class ObjectManager : public ObjectManagerInterface, ObjectID, std::unordered_map>> unfulfilled_push_requests_; - /// The objects that this object manager is currently trying to fetch from - /// remote object managers. - std::unordered_map pull_requests_; - /// Profiling events that are to be batched together and added to the profile /// table in the GCS. std::vector profile_events_; @@ -460,9 +442,6 @@ class ObjectManager : public ObjectManagerInterface, /// and rpc thread. std::mutex profile_mutex_; - /// Internally maintained random number generator. - std::mt19937_64 gen_; - /// The gPRC server. rpc::GrpcServer object_manager_server_; @@ -478,9 +457,16 @@ class ObjectManager : public ObjectManagerInterface, const RestoreSpilledObjectCallback restore_spilled_object_; + /// Pull manager retry timer . + /* std::unique_ptr pull_retry_timer_; */ + boost::asio::deadline_timer pull_retry_timer_; + /// Object push manager. std::unique_ptr push_manager_; + /// Object pull manager. + std::unique_ptr pull_manager_; + /// Running sum of the amount of memory used in the object store. int64_t used_memory_ = 0; diff --git a/src/ray/object_manager/pull_manager.cc b/src/ray/object_manager/pull_manager.cc new file mode 100644 index 000000000..082426cc1 --- /dev/null +++ b/src/ray/object_manager/pull_manager.cc @@ -0,0 +1,139 @@ +#include "ray/object_manager/pull_manager.h" + +namespace ray { + +PullManager::PullManager( + NodeID &self_node_id, const std::function object_is_local, + const std::function send_pull_request, + const RestoreSpilledObjectCallback restore_spilled_object, + const std::function get_time, int pull_timeout_ms) + : self_node_id_(self_node_id), + object_is_local_(object_is_local), + send_pull_request_(send_pull_request), + restore_spilled_object_(restore_spilled_object), + get_time_(get_time), + pull_timeout_ms_(pull_timeout_ms), + gen_(std::chrono::high_resolution_clock::now().time_since_epoch().count()) {} + +bool PullManager::Pull(const ObjectID &object_id, const rpc::Address &owner_address) { + RAY_LOG(DEBUG) << "Pull " + << " of object " << object_id; + // Check if object is already local. + if (object_is_local_(object_id)) { + RAY_LOG(DEBUG) << object_id << " attempted to pull an object that's already local."; + return false; + } + if (pull_requests_.find(object_id) != pull_requests_.end()) { + RAY_LOG(DEBUG) << object_id << " has inflight pull_requests, skipping."; + return false; + } + + pull_requests_.emplace(object_id, PullRequest(get_time_() + pull_timeout_ms_ / 1000)); + return true; +} + +void PullManager::OnLocationChange(const ObjectID &object_id, + const std::unordered_set &client_ids, + const std::string &spilled_url) { + // 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 = std::vector(client_ids.begin(), client_ids.end()); + if (!spilled_url.empty()) { + // Try to restore the spilled object. + restore_spilled_object_(object_id, spilled_url, + [this, object_id](const ray::Status &status) { + // Fall back to fetching from another object manager. + if (!status.ok()) { + TryPull(object_id); + } + }); + } else { + // New object locations were found, so begin trying to pull from a + // client. This will be called every time a new client location + // appears. + TryPull(object_id); + } +} + +void PullManager::TryPull(const ObjectID &object_id) { + auto it = pull_requests_.find(object_id); + if (it == pull_requests_.end()) { + return; + } + + auto &node_vector = it->second.client_locations; + + // The timer should never fire if there are no expected client locations. + if (node_vector.empty()) { + return; + } + + RAY_CHECK(!object_is_local_(object_id)); + // Make sure that there is at least one client which is not the local client. + // TODO(rkn): It may actually be possible for this check to fail. + if (node_vector.size() == 1 && node_vector[0] == self_node_id_) { + RAY_LOG(WARNING) << "The object manager with ID " << self_node_id_ + << " is trying to pull object " << object_id + << " but the object table suggests that this object manager " + << "already has the object. The object may have been evicted. It is " + << "most likely due to memory pressure, object pull has been " + << "requested before object location is updated."; + return; + } + + // Choose a random client to pull the object from. + // Generate a random index. + std::uniform_int_distribution distribution(0, node_vector.size() - 1); + int node_index = distribution(gen_); + NodeID node_id = node_vector[node_index]; + // If the object manager somehow ended up choosing itself, choose a different + // object manager. + if (node_id == self_node_id_) { + std::swap(node_vector[node_index], node_vector[node_vector.size() - 1]); + node_vector.pop_back(); + RAY_LOG(WARNING) + << "The object manager with ID " << self_node_id_ << " is trying to pull object " + << object_id << " but the object table suggests that this object manager " + << "already has the object. It is most likely due to memory pressure, object " + << "pull has been requested before object location is updated."; + node_id = node_vector[node_index % node_vector.size()]; + RAY_CHECK(node_id != self_node_id_); + } + + RAY_LOG(DEBUG) << "Sending pull request from " << self_node_id_ << " to " << node_id + << " of object " << object_id; + send_pull_request_(object_id, node_id); +} + +bool PullManager::CancelPull(const ObjectID &object_id) { + auto it = pull_requests_.find(object_id); + if (it == pull_requests_.end()) { + return false; + } + + pull_requests_.erase(it); + return true; +} + +void PullManager::Tick() { + for (auto &pair : pull_requests_) { + const auto &object_id = pair.first; + auto &request = pair.second; + const auto time = get_time_(); + if (time >= request.next_pull_time) { + TryPull(object_id); + request.next_pull_time = time + pull_timeout_ms_ / 1000; + } + } +} + +int PullManager::NumActiveRequests() const { return pull_requests_.size(); } + +} // namespace ray diff --git a/src/ray/object_manager/pull_manager.h b/src/ray/object_manager/pull_manager.h new file mode 100644 index 000000000..023f72d0e --- /dev/null +++ b/src/ray/object_manager/pull_manager.h @@ -0,0 +1,115 @@ + +#pragma once + +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/time/clock.h" +#include "ray/common/id.h" +#include "ray/common/ray_config.h" +#include "ray/common/status.h" +#include "ray/object_manager/common.h" +#include "ray/object_manager/format/object_manager_generated.h" +#include "ray/object_manager/notification/object_store_notification_manager_ipc.h" +#include "ray/object_manager/object_buffer_pool.h" +#include "ray/object_manager/object_directory.h" +#include "ray/object_manager/ownership_based_object_directory.h" +#include "ray/object_manager/plasma/store_runner.h" +#include "ray/rpc/object_manager/object_manager_client.h" +#include "ray/rpc/object_manager/object_manager_server.h" + +namespace ray { + +class PullManager { + public: + /// PullManager is responsible for managing the policy around when to send pull requests + /// and to whom. Notably, it is _not_ responsible for controlling the object directory + /// or any pubsub communications. + /// + /// \param self_node_id the current node + /// \param object_is_local A callback which should return true if a given object is + /// already on the local node. \param send_pull_request A callback which should send a + /// pull request to the specified node. + /// \param restore_spilled_object A callback which should + /// retrieve an spilled object from the external store. + PullManager( + NodeID &self_node_id, const std::function object_is_local, + const std::function send_pull_request, + const RestoreSpilledObjectCallback restore_spilled_object, + const std::function get_time, int pull_timeout_ms); + + /// Begin a new pull request if necessary. + /// + /// \param object_id The object id to pull. + /// \param owner_address The owner of the object. + /// + /// \return True if a new pull request was necessary. If true, the caller should + /// subscribe to new locations of the object, and call OnLocationChange when necessary. + bool Pull(const ObjectID &object_id, const rpc::Address &owner_address); + + /// Called when the available locations for a given object change. + /// + /// \param object_id The ID of the object which is now available in a new location. + /// \param client_ids The new set of nodes that the object is available on. Not + /// necessarily a super or subset of the previously available nodes. \param spilled_url + /// The location of the object if it was spilled. If non-empty, the object may no longer + /// be on any node. + void OnLocationChange(const ObjectID &object_id, + const std::unordered_set &client_ids, + const std::string &spilled_url); + + /// Cancel an existing pull request if necessary. + /// + /// \param object_id The object id that no longer needs to be pulled. + /// + /// \return True if a pull was cancelled. If there was no pending pull request for the + /// object this method may return false. + bool CancelPull(const ObjectID &object_id); + + /// Called when the retry timer fires. If this fires, the pull manager may try to pull + /// existing objects from other nodes if necessary. + void Tick(); + + /// The number of ongoing object pulls. + int NumActiveRequests() const; + + private: + /// A helper structure for tracking information about each ongoing object pull. + struct PullRequest { + PullRequest(double first_retry_time) + : client_locations(), next_pull_time(first_retry_time) {} + std::vector client_locations; + double next_pull_time; + }; + + /// See the constructor's arguments. + NodeID self_node_id_; + const std::function object_is_local_; + const std::function send_pull_request_; + const RestoreSpilledObjectCallback restore_spilled_object_; + const std::function get_time_; + int pull_timeout_ms_; + + /// The objects that this object manager is currently trying to fetch from + /// remote object managers. + std::unordered_map pull_requests_; + + /// Internally maintained random number generator. + std::mt19937_64 gen_; + + /// 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. + /// \return Void. + void TryPull(const ObjectID &object_id); +}; +} // namespace ray diff --git a/src/ray/object_manager/test/pull_manager_test.cc b/src/ray/object_manager/test/pull_manager_test.cc new file mode 100644 index 000000000..90f34048c --- /dev/null +++ b/src/ray/object_manager/test/pull_manager_test.cc @@ -0,0 +1,173 @@ + +#include "ray/object_manager/pull_manager.h" + +#include "gtest/gtest.h" +#include "ray/common/test_util.h" + +namespace ray { + +class PullManagerTest : public ::testing::Test { + public: + PullManagerTest() + : self_node_id_(NodeID::FromRandom()), + object_is_local_(false), + num_send_pull_request_calls_(0), + num_restore_spilled_object_calls_(0), + fake_time_(0), + pull_manager_(self_node_id_, + [this](const ObjectID &object_id) { return object_is_local_; }, + [this](const ObjectID &object_id, const NodeID &node_id) { + num_send_pull_request_calls_++; + }, + [this](const ObjectID &, const std::string &, + std::function) { + num_restore_spilled_object_calls_++; + }, + [this]() { return fake_time_; }, 10000) {} + + NodeID self_node_id_; + bool object_is_local_; + int num_send_pull_request_calls_; + int num_restore_spilled_object_calls_; + double fake_time_; + PullManager pull_manager_; +}; + +TEST_F(PullManagerTest, TestStaleSubscription) { + ObjectID obj1 = ObjectID::FromRandom(); + rpc::Address addr1; + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); + pull_manager_.Pull(obj1, addr1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 1); + + std::unordered_set client_ids; + pull_manager_.OnLocationChange(obj1, client_ids, ""); + + // There are no client ids to pull from. + ASSERT_EQ(num_send_pull_request_calls_, 0); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + + pull_manager_.CancelPull(obj1); + + ASSERT_EQ(num_send_pull_request_calls_, 0); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); + + client_ids.insert(NodeID::FromRandom()); + pull_manager_.OnLocationChange(obj1, client_ids, ""); + + // Now we're getting a notification about an object that was already cancelled. + ASSERT_EQ(num_send_pull_request_calls_, 0); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); +} + +TEST_F(PullManagerTest, TestRestoreSpilledObject) { + ObjectID obj1 = ObjectID::FromRandom(); + rpc::Address addr1; + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); + pull_manager_.Pull(obj1, addr1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 1); + + std::unordered_set client_ids; + pull_manager_.OnLocationChange(obj1, client_ids, "remote_url/foo/bar"); + + // client_ids is empty here, so there's nowhere to pull from. + ASSERT_EQ(num_send_pull_request_calls_, 0); + ASSERT_EQ(num_restore_spilled_object_calls_, 1); + + client_ids.insert(NodeID::FromRandom()); + pull_manager_.OnLocationChange(obj1, client_ids, "remote_url/foo/bar"); + + // The behavior is supposed to be to always restore the spilled object if possible (even + // if it exists elsewhere in the cluster). + ASSERT_EQ(num_send_pull_request_calls_, 0); + ASSERT_EQ(num_restore_spilled_object_calls_, 2); + + pull_manager_.CancelPull(obj1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); +} + +TEST_F(PullManagerTest, TestManyUpdates) { + ObjectID obj1 = ObjectID::FromRandom(); + rpc::Address addr1; + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); + pull_manager_.Pull(obj1, addr1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 1); + + std::unordered_set client_ids; + client_ids.insert(NodeID::FromRandom()); + + for (int i = 0; i < 100; i++) { + pull_manager_.OnLocationChange(obj1, client_ids, ""); + } + + ASSERT_EQ(num_send_pull_request_calls_, 100); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + + pull_manager_.CancelPull(obj1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); +} + +TEST_F(PullManagerTest, TestRetryTimer) { + ObjectID obj1 = ObjectID::FromRandom(); + rpc::Address addr1; + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); + pull_manager_.Pull(obj1, addr1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 1); + + std::unordered_set client_ids; + client_ids.insert(NodeID::FromRandom()); + + // We need to call OnLocationChange at least once, to population the list of nodes with + // the object. + pull_manager_.OnLocationChange(obj1, client_ids, ""); + ASSERT_EQ(num_send_pull_request_calls_, 1); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + + for (; fake_time_ <= 127 * 10; fake_time_ += 0.1) { + pull_manager_.Tick(); + } + + // Rapid set of location changes. + for (int i = 0; i < 127; i++) { + fake_time_ += 0.1; + pull_manager_.OnLocationChange(obj1, client_ids, ""); + } + + // We should make a pull request every tick (even if it's a duplicate to a node we're + // already pulling from). + // OnLocationChange also doesn't count towards the retry timer. + // To the casual observer, this may seem off-by-one, but this is due to floating point + // error (0.1 + 0.1 ... 10k times > 10 == True) + ASSERT_EQ(num_send_pull_request_calls_, 127 * 2); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + + pull_manager_.CancelPull(obj1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); +} + +TEST_F(PullManagerTest, TestBasic) { + ObjectID obj1 = ObjectID::FromRandom(); + rpc::Address addr1; + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); + pull_manager_.Pull(obj1, addr1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 1); + + std::unordered_set client_ids; + client_ids.insert(NodeID::FromRandom()); + pull_manager_.OnLocationChange(obj1, client_ids, ""); + + ASSERT_EQ(num_send_pull_request_calls_, 1); + ASSERT_EQ(num_restore_spilled_object_calls_, 0); + + pull_manager_.CancelPull(obj1); + ASSERT_EQ(pull_manager_.NumActiveRequests(), 0); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/ray/raylet/main.cc b/src/ray/raylet/main.cc index 4f80d5ca3..fdb858d51 100644 --- a/src/ray/raylet/main.cc +++ b/src/ray/raylet/main.cc @@ -234,6 +234,9 @@ int main(int argc, char *argv[]) { ray::ObjectManagerConfig object_manager_config; object_manager_config.object_manager_port = object_manager_port; object_manager_config.store_socket_name = store_socket_name; + + object_manager_config.timer_freq_ms = + RayConfig::instance().object_manager_timer_freq_ms(); object_manager_config.pull_timeout_ms = RayConfig::instance().object_manager_pull_timeout_ms(); object_manager_config.push_timeout_ms =