From 427b5af0aefafa573c690487a3fc3420e7b2816b Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Wed, 28 Oct 2020 10:38:42 -0400 Subject: [PATCH] [Object spilling] Refactor raylet to add a local object manager class (#11647) * Fix pytest... * Release objects that have been spilled * GCS object table interface refactor * Add spilled URL to object location info * refactor to include spilled URL in notifications * improve tests * Add spilled URL to object directory results * Remove force restore call * Merge spilled URL and location * fix * tmp * refactor * unit test skeleton * unit testing * unit test fixes * cleanup * cleanup * update * Separate pinning from waiting for object free, fixes pytest * Update src/ray/raylet/local_object_manager.h Co-authored-by: Eric Liang Co-authored-by: Tyler Westenbroek Co-authored-by: Eric Liang --- BUILD.bazel | 12 + src/ray/raylet/local_object_manager.cc | 178 ++++++++++++ src/ray/raylet/local_object_manager.h | 134 +++++++++ src/ray/raylet/main.cc | 2 - src/ray/raylet/node_manager.cc | 204 ++------------ src/ray/raylet/node_manager.h | 43 +-- src/ray/raylet/raylet.cc | 13 +- .../scheduling/cluster_task_manager_test.cc | 190 +------------ .../raylet/test/local_object_manager_test.cc | 259 ++++++++++++++++++ src/ray/raylet/test/util.h | 212 ++++++++++++++ src/ray/raylet/worker.h | 4 +- src/ray/raylet/worker_pool.h | 24 +- 12 files changed, 868 insertions(+), 407 deletions(-) create mode 100644 src/ray/raylet/local_object_manager.cc create mode 100644 src/ray/raylet/local_object_manager.h create mode 100644 src/ray/raylet/test/local_object_manager_test.cc create mode 100644 src/ray/raylet/test/util.h diff --git a/BUILD.bazel b/BUILD.bazel index 6494aa563..4d31f64f6 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -794,6 +794,18 @@ cc_test( ], ) +cc_test( + name = "local_object_manager_test", + srcs = [ + "src/ray/raylet/test/local_object_manager_test.cc", + ], + copts = COPTS, + deps = [ + ":raylet_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "reconstruction_policy_test", srcs = ["src/ray/raylet/reconstruction_policy_test.cc"], diff --git a/src/ray/raylet/local_object_manager.cc b/src/ray/raylet/local_object_manager.cc new file mode 100644 index 000000000..fce7536b3 --- /dev/null +++ b/src/ray/raylet/local_object_manager.cc @@ -0,0 +1,178 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ray/raylet/local_object_manager.h" + +namespace ray { + +namespace raylet { + +void LocalObjectManager::PinObjects(const std::vector &object_ids, + std::vector> &&objects) { + for (size_t i = 0; i < object_ids.size(); i++) { + const auto &object_id = object_ids[i]; + auto &object = objects[i]; + if (object == nullptr) { + RAY_LOG(ERROR) << "Plasma object " << object_id + << " was evicted before the raylet could pin it."; + continue; + } + RAY_LOG(DEBUG) << "Pinning object " << object_id; + pinned_objects_.emplace(object_id, std::move(object)); + } +} + +void LocalObjectManager::WaitForObjectFree(const rpc::Address &owner_address, + const std::vector &object_ids) { + for (const auto &object_id : object_ids) { + // Send a long-running RPC request to the owner for each object. When we get a + // response or the RPC fails (due to the owner crashing), unpin the object. + // TODO(edoakes): we should be batching these requests instead of sending one per + // pinned object. + rpc::WaitForObjectEvictionRequest wait_request; + wait_request.set_object_id(object_id.Binary()); + wait_request.set_intended_worker_id(owner_address.worker_id()); + auto owner_client = owner_client_pool_.GetOrConnect(owner_address); + owner_client->WaitForObjectEviction( + wait_request, + [this, object_id](Status status, const rpc::WaitForObjectEvictionReply &reply) { + if (!status.ok()) { + RAY_LOG(WARNING) << "Worker failed. Unpinning object " << object_id; + } + ReleaseFreedObject(object_id); + }); + } +} + +void LocalObjectManager::ReleaseFreedObject(const ObjectID &object_id) { + RAY_LOG(DEBUG) << "Unpinning object " << object_id; + pinned_objects_.erase(object_id); + + // Try to evict all copies of the object from the cluster. + if (free_objects_period_ms_ >= 0) { + objects_to_free_.push_back(object_id); + } + if (objects_to_free_.size() == free_objects_batch_size_ || + free_objects_period_ms_ == 0) { + FlushFreeObjects(); + } +} + +void LocalObjectManager::FlushFreeObjects() { + if (!objects_to_free_.empty()) { + RAY_LOG(DEBUG) << "Freeing " << objects_to_free_.size() << " out-of-scope objects"; + on_objects_freed_(objects_to_free_); + objects_to_free_.clear(); + } + last_free_objects_at_ms_ = current_time_ms(); +} + +void LocalObjectManager::FlushFreeObjectsIfNeeded(int64_t now_ms) { + if (free_objects_period_ms_ > 0 && + static_cast(now_ms - last_free_objects_at_ms_) > free_objects_period_ms_) { + FlushFreeObjects(); + } +} + +void LocalObjectManager::SpillObjects(const std::vector &object_ids, + std::function callback) { + for (const auto &id : object_ids) { + // We should not spill an object that we are not the primary copy for. + if (pinned_objects_.count(id) == 0) { + callback( + Status::Invalid("Requested spill for object that is not marked as " + "the primary copy.")); + } + } + + io_worker_pool_.PopIOWorker( + [this, object_ids, callback](std::shared_ptr io_worker) { + rpc::SpillObjectsRequest request; + for (const auto &object_id : object_ids) { + RAY_LOG(DEBUG) << "Sending spill request for object " << object_id; + request.add_object_ids_to_spill(object_id.Binary()); + } + io_worker->rpc_client()->SpillObjects( + request, [this, object_ids, callback, io_worker]( + const ray::Status &status, const rpc::SpillObjectsReply &r) { + io_worker_pool_.PushIOWorker(io_worker); + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to send object spilling request: " + << status.ToString(); + if (callback) { + callback(status); + } + } else { + AddSpilledUrls(object_ids, r, callback); + } + }); + }); +} + +void LocalObjectManager::AddSpilledUrls( + const std::vector &object_ids, const rpc::SpillObjectsReply &worker_reply, + std::function callback) { + auto num_remaining = std::make_shared(object_ids.size()); + for (size_t i = 0; i < object_ids.size(); ++i) { + const ObjectID &object_id = object_ids[i]; + const std::string &object_url = worker_reply.spilled_objects_url(i); + RAY_LOG(DEBUG) << "Object " << object_id << " spilled at " << object_url; + // Write to object directory. Wait for the write to finish before + // releasing the object to make sure that the spilled object can + // be retrieved by other raylets. + RAY_CHECK_OK(object_info_accessor_.AsyncAddSpilledUrl( + object_id, object_url, [this, object_id, callback, num_remaining](Status status) { + RAY_CHECK_OK(status); + // Unpin the object. + // NOTE(swang): Due to a race condition, the object may not be in + // the map yet. In that case, the owner will respond to the + // WaitForObjectEvictionRequest and we will unpin the object + // then. + pinned_objects_.erase(object_id); + (*num_remaining)--; + if (*num_remaining == 0 && callback) { + callback(status); + } + })); + } +} + +void LocalObjectManager::AsyncRestoreSpilledObject( + const ObjectID &object_id, const std::string &object_url, + std::function callback) { + RAY_LOG(DEBUG) << "Restoring spilled object " << object_id << " from URL " + << object_url; + io_worker_pool_.PopIOWorker([this, object_url, + callback](std::shared_ptr io_worker) { + RAY_LOG(DEBUG) << "Sending restore spilled object request"; + rpc::RestoreSpilledObjectsRequest request; + request.add_spilled_objects_url(std::move(object_url)); + io_worker->rpc_client()->RestoreSpilledObjects( + request, [this, callback, io_worker](const ray::Status &status, + const rpc::RestoreSpilledObjectsReply &r) { + io_worker_pool_.PushIOWorker(io_worker); + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to send restore spilled object request: " + << status.ToString(); + } + if (callback) { + callback(status); + } + }); + }); +} + +}; // namespace raylet + +}; // namespace ray diff --git a/src/ray/raylet/local_object_manager.h b/src/ray/raylet/local_object_manager.h new file mode 100644 index 000000000..719c03d15 --- /dev/null +++ b/src/ray/raylet/local_object_manager.h @@ -0,0 +1,134 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +#include + +#include "ray/common/id.h" +#include "ray/common/ray_object.h" +#include "ray/gcs/accessor.h" +#include "ray/raylet/worker_pool.h" +#include "ray/rpc/worker/core_worker_client_pool.h" + +namespace ray { + +namespace raylet { + +/// This class implements memory management for primary objects, objects that +/// have been freed, and objects that have been spilled. +class LocalObjectManager { + public: + LocalObjectManager(size_t free_objects_batch_size, int64_t free_objects_period_ms, + IOWorkerPoolInterface &io_worker_pool, + gcs::ObjectInfoAccessor &object_info_accessor, + rpc::CoreWorkerClientPool &owner_client_pool, + std::function &)> on_objects_freed) + : free_objects_period_ms_(free_objects_period_ms), + free_objects_batch_size_(free_objects_batch_size), + io_worker_pool_(io_worker_pool), + object_info_accessor_(object_info_accessor), + owner_client_pool_(owner_client_pool), + on_objects_freed_(on_objects_freed), + last_free_objects_at_ms_(current_time_ms()) {} + + /// Pin objects. + /// + /// \param object_ids The objects to be pinned. + /// \param objects Pointers to the objects to be pinned. The pointer should + /// be kept in scope until the object can be released. + void PinObjects(const std::vector &object_ids, + std::vector> &&objects); + + /// Wait for the objects' owner to free the object. The objects will be + /// released when the owner at the given address fails or replies that the + /// object can be evicted. + /// + /// \param owner_address The address of the owner of the objects. + /// \param object_ids The objects to be freed. + void WaitForObjectFree(const rpc::Address &owner_address, + const std::vector &object_ids); + + /// Spill objects to external storage. + /// + /// \param objects_ids_to_spill The objects to be spilled. + /// \param callback A callback to call once the objects have been spilled, or + /// there is an error. + void SpillObjects(const std::vector &objects_ids, + std::function callback); + + /// Restore a spilled object from external storage back into local memory. + /// + /// \param object_id The ID of the object to restore. + /// \param object_url The URL in external storage from which the object can be restored. + /// \param callback A callback to call when the restoration is done. Status + /// will contain the error during restoration, if any. + void AsyncRestoreSpilledObject(const ObjectID &object_id, const std::string &object_url, + std::function callback); + + /// Try to clear any objects that have been freed. + void FlushFreeObjectsIfNeeded(int64_t now_ms); + + private: + /// Release an object that has been freed by its owner. + void ReleaseFreedObject(const ObjectID &object_id); + + /// Clear any freed objects. This will trigger the callback for freed + /// objects. + void FlushFreeObjects(); + + /// Add objects' spilled URLs to the global object directory. Call the + /// callback once all URLs have been added. + void AddSpilledUrls(const std::vector &object_ids, + const rpc::SpillObjectsReply &worker_reply, + std::function callback); + + /// The period between attempts to eagerly evict objects from plasma. + const int64_t free_objects_period_ms_; + + /// The number of freed objects to accumulate before flushing. + const size_t free_objects_batch_size_; + + /// A worker pool, used for spilling and restoring objects. + IOWorkerPoolInterface &io_worker_pool_; + + /// A GCS client, used to update locations for spilled objects. + gcs::ObjectInfoAccessor &object_info_accessor_; + + /// Cache of gRPC clients to owners of objects pinned on + /// this node. + rpc::CoreWorkerClientPool &owner_client_pool_; + + /// A callback to call when an object has been freed. + std::function &)> on_objects_freed_; + + // Objects that are pinned on this node. + absl::flat_hash_map> pinned_objects_; + + /// The time that we last sent a FreeObjects request to other nodes for + /// objects that have gone out of scope in the application. + uint64_t last_free_objects_at_ms_ = 0; + + /// Objects that are out of scope in the application and that should be freed + /// from plasma. The cache is flushed when it reaches the + /// free_objects_batch_size, or if objects have been in the cache for longer + /// than the config's free_objects_period, whichever occurs first. + std::vector objects_to_free_; +}; + +}; // namespace raylet + +}; // namespace ray diff --git a/src/ray/raylet/main.cc b/src/ray/raylet/main.cc index 4ee33551b..816f52017 100644 --- a/src/ray/raylet/main.cc +++ b/src/ray/raylet/main.cc @@ -219,8 +219,6 @@ int main(int argc, char *argv[]) { RayConfig::instance().raylet_heartbeat_timeout_milliseconds(); node_manager_config.debug_dump_period_ms = RayConfig::instance().debug_dump_period_milliseconds(); - node_manager_config.free_objects_period_ms = - RayConfig::instance().free_objects_period_milliseconds(); node_manager_config.fair_queueing_enabled = RayConfig::instance().fair_queueing_enabled(); node_manager_config.object_pinning_enabled = diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 8b884587d..44226fc22 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -130,7 +130,6 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, const NodeID &self heartbeat_timer_(io_service), heartbeat_period_(std::chrono::milliseconds(config.heartbeat_period_ms)), debug_dump_period_(config.debug_dump_period_ms), - free_objects_period_(config.free_objects_period_ms), fair_queueing_enabled_(config.fair_queueing_enabled), object_pinning_enabled_(config.object_pinning_enabled), temp_dir_(config.temp_dir), @@ -162,6 +161,14 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, const NodeID &self new DefaultAgentManagerServiceHandler(agent_manager_)), agent_manager_service_(io_service, *agent_manager_service_handler_), client_call_manager_(io_service), + worker_rpc_pool_(client_call_manager_), + local_object_manager_(RayConfig::instance().free_objects_batch_size(), + RayConfig::instance().free_objects_period_milliseconds(), + worker_pool_, gcs_client_->Objects(), worker_rpc_pool_, + [this](const std::vector &object_ids) { + object_manager_.FreeObjects(object_ids, + /*local_only=*/false); + }), new_scheduler_enabled_(RayConfig::instance().new_scheduler_enabled()), report_worker_backlog_(RayConfig::instance().report_worker_backlog()) { RAY_LOG(INFO) << "Initializing NodeManager with ID " << self_node_id_; @@ -302,7 +309,6 @@ ray::Status NodeManager::RegisterGcs() { // Start sending heartbeats to the GCS. last_heartbeat_at_ms_ = current_time_ms(); last_debug_dump_at_ms_ = current_time_ms(); - last_free_objects_at_ms_ = current_time_ms(); Heartbeat(); // Start the timer that gets object manager profiling information and sends it // to the GCS. @@ -495,10 +501,7 @@ void NodeManager::Heartbeat() { } // Evict all copies of freed objects from the cluster. - if (free_objects_period_ > 0 && - static_cast(now_ms - last_free_objects_at_ms_) > free_objects_period_) { - FlushObjectsToFree(); - } + local_object_manager_.FlushFreeObjectsIfNeeded(now_ms); // Reset the timer. heartbeat_timer_.expires_from_now(heartbeat_period_); @@ -529,95 +532,14 @@ void NodeManager::DoLocalGC() { void NodeManager::HandleRequestObjectSpillage( const rpc::RequestObjectSpillageRequest &request, rpc::RequestObjectSpillageReply *reply, rpc::SendReplyCallback send_reply_callback) { - SpillObjects({ObjectID::FromBinary(request.object_id())}, - [reply, send_reply_callback](const ray::Status &status) { - if (status.ok()) { - reply->set_success(true); - } - send_reply_callback(Status::OK(), nullptr, nullptr); - }); -} - -void NodeManager::SpillObjects(const std::vector &objects_ids, - std::function callback) { - for (const auto &id : objects_ids) { - // We should not spill an object that we are not the primary copy for. - // TODO(swang): We should really return an error here but right now there - // is a race condition where the raylet receives the owner's request to - // spill an object before it receives the message to pin the objects from - // the local worker. - if (pinned_objects_.count(id) == 0) { - RAY_LOG(WARNING) << "Requested spill for object that has not yet been marked as " - "the primary copy."; - } - } - worker_pool_.PopIOWorker([this, objects_ids, - callback](std::shared_ptr io_worker) { - rpc::SpillObjectsRequest request; - for (const auto &object_id : objects_ids) { - RAY_LOG(DEBUG) << "Sending spill request for object " << object_id; - request.add_object_ids_to_spill(object_id.Binary()); - } - io_worker->rpc_client()->SpillObjects( - request, [this, objects_ids, callback, io_worker]( - const ray::Status &status, const rpc::SpillObjectsReply &r) { - worker_pool_.PushIOWorker(io_worker); - if (!status.ok()) { - RAY_LOG(ERROR) << "Failed to send object spilling request: " - << status.ToString(); - if (callback) { - callback(status); - } - } else { - for (size_t i = 0; i < objects_ids.size(); ++i) { - const ObjectID &object_id = objects_ids[i]; - const std::string &object_url = r.spilled_objects_url(i); - RAY_LOG(DEBUG) << "Object " << object_id << " spilled at " << object_url; - // Write to object directory. Wait for the write to finish before - // releasing the object to make sure that the spilled object can - // be retrieved by other raylets. - RAY_CHECK_OK(gcs_client_->Objects().AsyncAddSpilledUrl( - object_id, object_url, [this, object_id, callback](Status status) { - RAY_CHECK_OK(status); - // Unpin the object. - // NOTE(swang): Due to a race condition, the object may not be in - // the map yet. In that case, the owner will respond to the - // WaitForObjectEvictionRequest and we will unpin the object - // then. - pinned_objects_.erase(object_id); - if (callback) { - callback(status); - } - })); - } - } - }); - }); -} - -void NodeManager::AsyncRestoreSpilledObject( - const ObjectID &object_id, const std::string &object_url, - std::function callback) { - RAY_LOG(DEBUG) << "Restoring spilled object " << object_id << " from URL " - << object_url; - worker_pool_.PopIOWorker([this, object_url, - callback](std::shared_ptr io_worker) { - RAY_LOG(DEBUG) << "Sending restore spilled object request"; - rpc::RestoreSpilledObjectsRequest request; - request.add_spilled_objects_url(std::move(object_url)); - io_worker->rpc_client()->RestoreSpilledObjects( - request, [this, callback, io_worker](const ray::Status &status, - const rpc::RestoreSpilledObjectsReply &r) { - worker_pool_.PushIOWorker(io_worker); - if (!status.ok()) { - RAY_LOG(ERROR) << "Failed to send restore spilled object request: " - << status.ToString(); - } - if (callback) { - callback(status); - } - }); - }); + local_object_manager_.SpillObjects( + {ObjectID::FromBinary(request.object_id())}, + [reply, send_reply_callback](const ray::Status &status) { + if (status.ok()) { + reply->set_success(true); + } + send_reply_callback(Status::OK(), nullptr, nullptr); + }); } // TODO(edoakes): this function is problematic because it both sends warnings spuriously @@ -3065,28 +2987,16 @@ std::string compact_tag_string(const opencensus::stats::ViewDescriptor &view, void NodeManager::HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request, rpc::PinObjectIDsReply *reply, rpc::SendReplyCallback send_reply_callback) { - WorkerID worker_id = WorkerID::FromBinary(request.owner_address().worker_id()); - auto it = worker_rpc_clients_.find(worker_id); - if (it == worker_rpc_clients_.end()) { - auto client = std::unique_ptr( - new rpc::CoreWorkerClient(request.owner_address(), client_call_manager_)); - it = worker_rpc_clients_ - .emplace(worker_id, - std::make_pair, size_t>( - std::move(client), 0)) - .first; + std::vector object_ids; + object_ids.reserve(request.object_ids_size()); + for (const auto &object_id_binary : request.object_ids()) { + object_ids.push_back(ObjectID::FromBinary(object_id_binary)); } - if (object_pinning_enabled_) { // Pin the objects in plasma by getting them and holding a reference to // the returned buffer. // NOTE: the caller must ensure that the objects already exist in plasma before // sending a PinObjectIDs request. - std::vector object_ids; - object_ids.reserve(request.object_ids_size()); - for (const auto &object_id_binary : request.object_ids()) { - object_ids.push_back(ObjectID::FromBinary(object_id_binary)); - } std::vector plasma_results; // TODO(swang): This `Get` has a timeout of 0, so the plasma store will not // block when serving the request. However, if the plasma store is under @@ -3100,77 +3010,23 @@ void NodeManager::HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request, return; } - // Pin the requested objects until the owner notifies us that the objects can be - // unpinned by responding to the WaitForObjectEviction message. - // TODO(edoakes): we should be batching these requests instead of sending one per - // pinned object. + std::vector> objects; for (int64_t i = 0; i < request.object_ids().size(); i++) { - ObjectID object_id = ObjectID::FromBinary(request.object_ids(i)); - if (plasma_results[i].data == nullptr) { - RAY_LOG(ERROR) << "Plasma object " << object_id - << " was evicted before the raylet could pin it."; - continue; + objects.push_back(nullptr); + } else { + objects.emplace_back(std::unique_ptr(new RayObject( + std::make_shared(plasma_results[i].data), + std::make_shared(plasma_results[i].metadata), {}))); } - - RAY_LOG(DEBUG) << "Pinning object " << object_id; - RAY_CHECK( - pinned_objects_ - .emplace( - object_id, - std::unique_ptr(new RayObject( - std::make_shared(plasma_results[i].data), - std::make_shared(plasma_results[i].metadata), {}))) - .second); } + local_object_manager_.PinObjects(object_ids, std::move(objects)); } - - for (const auto &object_id_binary : request.object_ids()) { - ObjectID object_id = ObjectID::FromBinary(object_id_binary); - // Send a long-running RPC request to the owner for each object. When we get a - // response or the RPC fails (due to the owner crashing), unpin the object. - rpc::WaitForObjectEvictionRequest wait_request; - wait_request.set_object_id(object_id_binary); - wait_request.set_intended_worker_id(request.owner_address().worker_id()); - worker_rpc_clients_[worker_id].second++; - it->second.first->WaitForObjectEviction( - wait_request, [this, worker_id, object_id]( - Status status, const rpc::WaitForObjectEvictionReply &reply) { - if (!status.ok()) { - RAY_LOG(WARNING) << "Worker " << worker_id << " failed. Unpinning object " - << object_id; - } - RAY_LOG(DEBUG) << "Unpinning object " << object_id; - pinned_objects_.erase(object_id); - - // Try to evict all copies of the object from the cluster. - if (free_objects_period_ >= 0) { - objects_to_free_.push_back(object_id); - } - if (objects_to_free_.size() == - RayConfig::instance().free_objects_batch_size() || - free_objects_period_ == 0) { - FlushObjectsToFree(); - } - - // Remove the cached worker client if there are no more pending requests. - if (--worker_rpc_clients_[worker_id].second == 0) { - worker_rpc_clients_.erase(worker_id); - } - }); - } + // Wait for the object to be freed by the owner, which keeps the ref count. + local_object_manager_.WaitForObjectFree(request.owner_address(), object_ids); send_reply_callback(Status::OK(), nullptr, nullptr); } -void NodeManager::FlushObjectsToFree() { - if (!objects_to_free_.empty()) { - RAY_LOG(DEBUG) << "Freeing " << objects_to_free_.size() << " out-of-scope objects"; - object_manager_.FreeObjects(objects_to_free_, /*local_only=*/false); - objects_to_free_.clear(); - } - last_free_objects_at_ms_ = current_time_ms(); -} - void NodeManager::HandleGetNodeStats(const rpc::GetNodeStatsRequest &node_stats_request, rpc::GetNodeStatsReply *reply, rpc::SendReplyCallback send_reply_callback) { diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 0d8836c70..57b26a432 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -28,6 +28,7 @@ #include "ray/object_manager/object_manager.h" #include "ray/raylet/actor_registration.h" #include "ray/raylet/agent_manager.h" +#include "ray/raylet/local_object_manager.h" #include "ray/raylet/scheduling/scheduling_ids.h" #include "ray/raylet/scheduling/cluster_resource_scheduler.h" #include "ray/raylet/scheduling/cluster_task_manager.h" @@ -36,6 +37,7 @@ #include "ray/raylet/reconstruction_policy.h" #include "ray/raylet/task_dependency_manager.h" #include "ray/raylet/worker_pool.h" +#include "ray/rpc/worker/core_worker_client_pool.h" #include "ray/util/ordered_set.h" #include "ray/common/bundle_spec.h" // clang-format on @@ -85,9 +87,6 @@ struct NodeManagerConfig { uint64_t heartbeat_period_ms; /// The time between debug dumps in milliseconds, or -1 to disable. uint64_t debug_dump_period_ms; - /// The time between attempts to eagerly evict objects from plasma in - /// milliseconds, or -1 to disable. - int64_t free_objects_period_ms; /// Whether to enable fair queueing between task classes in raylet. bool fair_queueing_enabled; /// Whether to enable pinning for plasma objects. @@ -171,13 +170,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// Get the port of the node manager rpc server. int GetServerPort() const { return node_manager_server_.GetPort(); } - /// Restore a spilled object from external storage back into local memory. - /// \param object_id The ID of the object to restore. - /// \param object_url The URL in external storage from which the object can be restored. - /// \param callback A callback to call when the restoration is done. Status - /// will contain the error during restoration, if any. - void AsyncRestoreSpilledObject(const ObjectID &object_id, const std::string &object_url, - std::function callback); + LocalObjectManager &GetLocalObjectManager() { return local_object_manager_; } private: /// Methods for handling clients. @@ -656,11 +649,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// Trigger local GC on each worker of this raylet. void DoLocalGC(); - /// Spill objects to external storage. - /// \param objects_ids_to_spill The objects to be spilled. - void SpillObjects(const std::vector &objects_ids_to_spill, - std::function callback = nullptr); - /// Push an error to the driver if this node is full of actors and so we are /// unable to schedule new tasks or actors at all. void WarnResourceDeadlock(); @@ -702,8 +690,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { std::chrono::milliseconds heartbeat_period_; /// The period between debug state dumps. int64_t debug_dump_period_; - /// The period between attempts to eagerly evict objects from plasma. - int64_t free_objects_period_; /// Whether to enable fair queueing between task classes in raylet. bool fair_queueing_enabled_; /// Whether to enable pinning for plasma objects. @@ -727,9 +713,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { SchedulingResources last_heartbeat_resources_; /// The time that the last debug string was logged to the console. uint64_t last_debug_dump_at_ms_; - /// The time that we last sent a FreeObjects request to other nodes for - /// objects that have gone out of scope in the application. - uint64_t last_free_objects_at_ms_; /// The number of heartbeats that we should wait before sending the /// next load report. uint8_t num_heartbeats_before_load_report_; @@ -772,6 +755,13 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// as well as all `CoreWorkerClient`s. rpc::ClientCallManager client_call_manager_; + /// Pool of RPC client connections to core workers. + rpc::CoreWorkerClientPool worker_rpc_pool_; + + /// Manages all local objects that are pinned (primary + /// copies), freed, and/or spilled. + LocalObjectManager local_object_manager_; + /// Map from node ids to clients of the remote node managers. std::unordered_map> remote_node_manager_clients_; @@ -804,13 +794,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { std::shared_ptr new_resource_scheduler_; std::shared_ptr cluster_task_manager_; - /// Cache of gRPC clients to workers (not necessarily running on this node). - /// Also includes the number of inflight requests to each worker - when this - /// reaches zero, the client will be deleted and a new one will need to be created - /// for any subsequent requests. - absl::flat_hash_map, size_t>> - worker_rpc_clients_; - absl::flat_hash_map> pinned_objects_; // TODO(swang): Evict entries from these caches. @@ -826,12 +809,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { absl::flat_hash_map>> async_plasma_objects_notification_ GUARDED_BY(plasma_object_notification_lock_); - /// Objects that are out of scope in the application and that should be freed - /// from plasma. The cache is flushed when it reaches the config's - /// free_objects_batch_size, or if objects have been in the cache for longer - /// than the config's free_objects_period, whichever occurs first. - std::vector objects_to_free_; - /// This map represents the commit state of 2PC protocol for atomic placement group /// creation. absl::flat_hash_map, pair_hash> diff --git a/src/ray/raylet/raylet.cc b/src/ray/raylet/raylet.cc index 32dbab8aa..80f8aab27 100644 --- a/src/ray/raylet/raylet.cc +++ b/src/ray/raylet/raylet.cc @@ -69,12 +69,13 @@ Raylet::Raylet(boost::asio::io_service &main_service, const std::string &socket_ gcs_client_)) : std::dynamic_pointer_cast( std::make_shared(main_service, gcs_client_))), - object_manager_( - main_service, self_node_id_, object_manager_config, object_directory_, - [this](const ObjectID &object_id, const std::string &spilled_url, - std::function callback) { - node_manager_.AsyncRestoreSpilledObject(object_id, spilled_url, callback); - }), + object_manager_(main_service, self_node_id_, object_manager_config, + object_directory_, + [this](const ObjectID &object_id, const std::string &spilled_url, + std::function callback) { + node_manager_.GetLocalObjectManager().AsyncRestoreSpilledObject( + object_id, spilled_url, callback); + }), node_manager_(main_service, self_node_id_, node_manager_config, object_manager_, gcs_client_, object_directory_), socket_name_(socket_name), diff --git a/src/ray/raylet/scheduling/cluster_task_manager_test.cc b/src/ray/raylet/scheduling/cluster_task_manager_test.cc index 39a1f7217..4887a198d 100644 --- a/src/ray/raylet/scheduling/cluster_task_manager_test.cc +++ b/src/ray/raylet/scheduling/cluster_task_manager_test.cc @@ -27,6 +27,7 @@ #include "ray/common/test_util.h" #include "ray/raylet/scheduling/cluster_resource_scheduler.h" #include "ray/raylet/scheduling/scheduling_ids.h" +#include "ray/raylet/test/util.h" #ifdef UNORDERED_VS_ABSL_MAPS_EVALUATION #include @@ -56,195 +57,6 @@ class MockWorkerPool : public WorkerPoolInterface { std::list> workers; }; -class MockWorker : public WorkerInterface { - public: - MockWorker(WorkerID worker_id, int port) : worker_id_(worker_id), port_(port) {} - - WorkerID WorkerId() const { return worker_id_; } - - rpc::WorkerType GetWorkerType() const { return rpc::WorkerType::WORKER; } - - int Port() const { return port_; } - - void SetOwnerAddress(const rpc::Address &address) { address_ = address; } - - void AssignTaskId(const TaskID &task_id) {} - - void AssignJobId(const JobID &job_id) {} - - void SetAssignedTask(Task &assigned_task) {} - - const std::string IpAddress() const { return address_.ip_address(); } - - void SetAllocatedInstances( - std::shared_ptr &allocated_instances) { - allocated_instances_ = allocated_instances; - } - - void SetLifetimeAllocatedInstances( - std::shared_ptr &allocated_instances) { - lifetime_allocated_instances_ = allocated_instances; - } - - std::shared_ptr GetAllocatedInstances() { - return allocated_instances_; - } - std::shared_ptr GetLifetimeAllocatedInstances() { - return lifetime_allocated_instances_; - } - - void MarkDead() { RAY_CHECK(false) << "Method unused"; } - bool IsDead() const { - RAY_CHECK(false) << "Method unused"; - return false; - } - void MarkBlocked() { RAY_CHECK(false) << "Method unused"; } - void MarkUnblocked() { RAY_CHECK(false) << "Method unused"; } - bool IsBlocked() const { - RAY_CHECK(false) << "Method unused"; - return false; - } - - Process GetProcess() const { - RAY_CHECK(false) << "Method unused"; - return Process::CreateNewDummy(); - } - void SetProcess(Process proc) { RAY_CHECK(false) << "Method unused"; } - Language GetLanguage() const { - RAY_CHECK(false) << "Method unused"; - return Language::PYTHON; - } - - void Connect(int port) { RAY_CHECK(false) << "Method unused"; } - - int AssignedPort() const { - RAY_CHECK(false) << "Method unused"; - return -1; - } - void SetAssignedPort(int port) { RAY_CHECK(false) << "Method unused"; } - const TaskID &GetAssignedTaskId() const { - RAY_CHECK(false) << "Method unused"; - return TaskID::Nil(); - } - bool AddBlockedTaskId(const TaskID &task_id) { - RAY_CHECK(false) << "Method unused"; - return false; - } - bool RemoveBlockedTaskId(const TaskID &task_id) { - RAY_CHECK(false) << "Method unused"; - return false; - } - const std::unordered_set &GetBlockedTaskIds() const { - RAY_CHECK(false) << "Method unused"; - auto *t = new std::unordered_set(); - return *t; - } - const JobID &GetAssignedJobId() const { - RAY_CHECK(false) << "Method unused"; - return JobID::Nil(); - } - void AssignActorId(const ActorID &actor_id) { RAY_CHECK(false) << "Method unused"; } - const ActorID &GetActorId() const { - RAY_CHECK(false) << "Method unused"; - return ActorID::Nil(); - } - void MarkDetachedActor() { RAY_CHECK(false) << "Method unused"; } - bool IsDetachedActor() const { - RAY_CHECK(false) << "Method unused"; - return false; - } - const std::shared_ptr Connection() const { - RAY_CHECK(false) << "Method unused"; - return nullptr; - } - const rpc::Address &GetOwnerAddress() const { - RAY_CHECK(false) << "Method unused"; - return address_; - } - - const ResourceIdSet &GetLifetimeResourceIds() const { - RAY_CHECK(false) << "Method unused"; - auto *t = new ResourceIdSet(); - return *t; - } - void SetLifetimeResourceIds(ResourceIdSet &resource_ids) { - RAY_CHECK(false) << "Method unused"; - } - void ResetLifetimeResourceIds() { RAY_CHECK(false) << "Method unused"; } - - const ResourceIdSet &GetTaskResourceIds() const { - RAY_CHECK(false) << "Method unused"; - auto *t = new ResourceIdSet(); - return *t; - } - void SetTaskResourceIds(ResourceIdSet &resource_ids) { - RAY_CHECK(false) << "Method unused"; - } - void ResetTaskResourceIds() { RAY_CHECK(false) << "Method unused"; } - ResourceIdSet ReleaseTaskCpuResources() { - RAY_CHECK(false) << "Method unused"; - auto *t = new ResourceIdSet(); - return *t; - } - void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) { - RAY_CHECK(false) << "Method unused"; - } - - Status AssignTask(const Task &task, const ResourceIdSet &resource_id_set) { - RAY_CHECK(false) << "Method unused"; - Status s; - return s; - } - void DirectActorCallArgWaitComplete(int64_t tag) { - RAY_CHECK(false) << "Method unused"; - } - - void ClearAllocatedInstances() { allocated_instances_ = nullptr; } - - void ClearLifetimeAllocatedInstances() { RAY_CHECK(false) << "Method unused"; } - - void SetBorrowedCPUInstances(std::vector &cpu_instances) { - borrowed_cpu_instances_ = cpu_instances; - } - - const PlacementGroupID &GetPlacementGroupId() const { - RAY_CHECK(false) << "Method unused"; - return PlacementGroupID::Nil(); - } - - void SetPlacementGroupId(const PlacementGroupID &placement_group_id) { - RAY_CHECK(false) << "Method unused"; - } - - std::vector &GetBorrowedCPUInstances() { return borrowed_cpu_instances_; } - - void ClearBorrowedCPUInstances() { RAY_CHECK(false) << "Method unused"; } - - Task &GetAssignedTask() { - RAY_CHECK(false) << "Method unused"; - auto *t = new Task(); - return *t; - } - - bool IsRegistered() { - RAY_CHECK(false) << "Method unused"; - return false; - } - - rpc::CoreWorkerClient *rpc_client() { - RAY_CHECK(false) << "Method unused"; - return nullptr; - } - - private: - WorkerID worker_id_; - int port_; - rpc::Address address_; - std::shared_ptr allocated_instances_; - std::shared_ptr lifetime_allocated_instances_; - std::vector borrowed_cpu_instances_; -}; - std::shared_ptr CreateSingleNodeScheduler( const std::string &id) { std::unordered_map local_node_resources; diff --git a/src/ray/raylet/test/local_object_manager_test.cc b/src/ray/raylet/test/local_object_manager_test.cc new file mode 100644 index 000000000..9bdb2f9fc --- /dev/null +++ b/src/ray/raylet/test/local_object_manager_test.cc @@ -0,0 +1,259 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ray/raylet/local_object_manager.h" + +#include "gmock/gmock.h" +#include "gtest/gtest.h" +#include "ray/common/id.h" +#include "ray/gcs/accessor.h" +#include "ray/raylet/test/util.h" +#include "ray/raylet/worker_pool.h" +#include "ray/rpc/grpc_client.h" +#include "ray/rpc/worker/core_worker_client.h" +#include "ray/rpc/worker/core_worker_client_pool.h" +#include "src/ray/protobuf/core_worker.grpc.pb.h" +#include "src/ray/protobuf/core_worker.pb.h" + +namespace ray { + +namespace raylet { + +using ::testing::_; + +class MockWorkerClient : public rpc::CoreWorkerClientInterface { + public: + void WaitForObjectEviction( + const rpc::WaitForObjectEvictionRequest &request, + const rpc::ClientCallback &callback) override { + callbacks.push_back(callback); + } + + bool ReplyObjectEviction(Status status = Status::OK()) { + if (callbacks.size() == 0) { + return false; + } + auto callback = callbacks.front(); + auto reply = rpc::WaitForObjectEvictionReply(); + callback(status, reply); + callbacks.pop_front(); + return true; + } + + std::list> callbacks; +}; + +class MockIOWorkerClient : public rpc::CoreWorkerClientInterface { + public: + void SpillObjects( + const rpc::SpillObjectsRequest &request, + const rpc::ClientCallback &callback) override { + callbacks.push_back(callback); + } + + bool ReplySpillObjects(std::vector urls, Status status = Status::OK()) { + if (callbacks.size() == 0) { + return false; + } + auto callback = callbacks.front(); + auto reply = rpc::SpillObjectsReply(); + for (const auto &url : urls) { + reply.add_spilled_objects_url(url); + } + callback(status, reply); + callbacks.pop_front(); + return true; + } + + void RestoreSpilledObjects( + const rpc::RestoreSpilledObjectsRequest &request, + const rpc::ClientCallback &callback) override { + rpc::RestoreSpilledObjectsReply reply; + callback(Status(), reply); + } + + std::list> callbacks; +}; + +class MockIOWorker : public MockWorker { + public: + MockIOWorker(WorkerID worker_id, int port, + std::shared_ptr io_worker) + : MockWorker(worker_id, port), io_worker(io_worker) {} + + rpc::CoreWorkerClientInterface *rpc_client() { return io_worker.get(); } + + std::shared_ptr io_worker; +}; + +class MockIOWorkerPool : public IOWorkerPoolInterface { + public: + MOCK_METHOD1(PushIOWorker, void(const std::shared_ptr &worker)); + + void PopIOWorker( + std::function)> callback) override { + callback(io_worker); + } + + std::shared_ptr io_worker_client = + std::make_shared(); + std::shared_ptr io_worker = + std::make_shared(WorkerID::FromRandom(), 1234, io_worker_client); +}; + +class MockObjectInfoAccessor : public gcs::ObjectInfoAccessor { + public: + MOCK_METHOD2( + AsyncGetLocations, + Status(const ObjectID &object_id, + const gcs::OptionalItemCallback &callback)); + + MOCK_METHOD1(AsyncGetAll, + Status(const gcs::MultiItemCallback &callback)); + + MOCK_METHOD3(AsyncAddLocation, Status(const ObjectID &object_id, const NodeID &node_id, + const gcs::StatusCallback &callback)); + + Status AsyncAddSpilledUrl(const ObjectID &object_id, const std::string &spilled_url, + const gcs::StatusCallback &callback) { + object_urls[object_id] = spilled_url; + callback(Status()); + return Status(); + } + + MOCK_METHOD3(AsyncRemoveLocation, + Status(const ObjectID &object_id, const NodeID &node_id, + const gcs::StatusCallback &callback)); + + MOCK_METHOD3(AsyncSubscribeToLocations, + Status(const ObjectID &object_id, + const gcs::SubscribeCallback< + ObjectID, std::vector> &subscribe, + const gcs::StatusCallback &done)); + + MOCK_METHOD1(AsyncUnsubscribeToLocations, Status(const ObjectID &object_id)); + + MOCK_METHOD1(AsyncResubscribe, void(bool is_pubsub_server_restarted)); + + MOCK_METHOD1(IsObjectUnsubscribed, bool(const ObjectID &object_id)); + + std::unordered_map object_urls; +}; + +class LocalObjectManagerTest : public ::testing::Test { + public: + LocalObjectManagerTest() + : owner_client(std::make_shared()), + client_pool([&](const rpc::Address &addr) { return owner_client; }), + manager(free_objects_batch_size, + /*free_objects_period_ms=*/1000, worker_pool, object_table, client_pool, + [&](const std::vector &object_ids) { + for (const auto &object_id : object_ids) { + freed.insert(object_id); + } + }) {} + + size_t free_objects_batch_size = 3; + std::shared_ptr owner_client; + rpc::CoreWorkerClientPool client_pool; + MockIOWorkerPool worker_pool; + MockObjectInfoAccessor object_table; + LocalObjectManager manager; + + std::unordered_set freed; +}; + +TEST_F(LocalObjectManagerTest, TestPin) { + rpc::Address owner_address; + owner_address.set_worker_id(WorkerID::FromRandom().Binary()); + + std::vector object_ids; + std::vector> objects; + + for (size_t i = 0; i < free_objects_batch_size; i++) { + ObjectID object_id = ObjectID::FromRandom(); + object_ids.push_back(object_id); + std::string meta = std::to_string(static_cast(rpc::ErrorType::OBJECT_IN_PLASMA)); + auto metadata = const_cast(reinterpret_cast(meta.data())); + auto meta_buffer = std::make_shared(metadata, meta.size()); + std::unique_ptr object( + new RayObject(nullptr, meta_buffer, std::vector())); + objects.push_back(std::move(object)); + } + manager.PinObjects(object_ids, std::move(objects)); + manager.WaitForObjectFree(owner_address, object_ids); + + for (size_t i = 0; i < free_objects_batch_size; i++) { + ASSERT_TRUE(freed.empty()); + ASSERT_TRUE(owner_client->ReplyObjectEviction()); + } + std::unordered_set expected(object_ids.begin(), object_ids.end()); + ASSERT_EQ(freed, expected); +} + +TEST_F(LocalObjectManagerTest, TestRestoreSpilledObject) { + ObjectID object_id = ObjectID::FromRandom(); + std::string object_url("url"); + int num_times_fired = 0; + EXPECT_CALL(worker_pool, PushIOWorker(_)); + manager.AsyncRestoreSpilledObject(object_id, object_url, [&](const Status &status) { + ASSERT_TRUE(status.ok()); + num_times_fired++; + }); + ASSERT_EQ(num_times_fired, 1); +} + +TEST_F(LocalObjectManagerTest, TestExplicitSpill) { + std::vector object_ids; + std::vector> objects; + + for (size_t i = 0; i < free_objects_batch_size; i++) { + ObjectID object_id = ObjectID::FromRandom(); + object_ids.push_back(object_id); + std::string meta = std::to_string(static_cast(rpc::ErrorType::OBJECT_IN_PLASMA)); + auto metadata = const_cast(reinterpret_cast(meta.data())); + auto meta_buffer = std::make_shared(metadata, meta.size()); + std::unique_ptr object( + new RayObject(nullptr, meta_buffer, std::vector())); + objects.push_back(std::move(object)); + } + manager.PinObjects(object_ids, std::move(objects)); + + int num_times_fired = 0; + manager.SpillObjects(object_ids, [&](const Status &status) mutable { + ASSERT_TRUE(status.ok()); + num_times_fired++; + }); + ASSERT_EQ(num_times_fired, 0); + + EXPECT_CALL(worker_pool, PushIOWorker(_)); + std::vector urls; + for (size_t i = 0; i < object_ids.size(); i++) { + urls.push_back("url" + std::to_string(i)); + } + ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls)); + ASSERT_EQ(num_times_fired, 1); + for (size_t i = 0; i < object_ids.size(); i++) { + ASSERT_EQ(object_table.object_urls[object_ids[i]], urls[i]); + } +} + +} // namespace raylet + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/ray/raylet/test/util.h b/src/ray/raylet/test/util.h new file mode 100644 index 000000000..9339455c2 --- /dev/null +++ b/src/ray/raylet/test/util.h @@ -0,0 +1,212 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ray/raylet/worker.h" + +namespace ray { + +namespace raylet { + +class MockWorker : public WorkerInterface { + public: + MockWorker(WorkerID worker_id, int port) : worker_id_(worker_id), port_(port) {} + + WorkerID WorkerId() const { return worker_id_; } + + rpc::WorkerType GetWorkerType() const { return rpc::WorkerType::WORKER; } + + int Port() const { return port_; } + + void SetOwnerAddress(const rpc::Address &address) { address_ = address; } + + void AssignTaskId(const TaskID &task_id) {} + + void AssignJobId(const JobID &job_id) {} + + void SetAssignedTask(Task &assigned_task) {} + + const std::string IpAddress() const { return address_.ip_address(); } + + void SetAllocatedInstances( + std::shared_ptr &allocated_instances) { + allocated_instances_ = allocated_instances; + } + + void SetLifetimeAllocatedInstances( + std::shared_ptr &allocated_instances) { + lifetime_allocated_instances_ = allocated_instances; + } + + std::shared_ptr GetAllocatedInstances() { + return allocated_instances_; + } + std::shared_ptr GetLifetimeAllocatedInstances() { + return lifetime_allocated_instances_; + } + + void MarkDead() { RAY_CHECK(false) << "Method unused"; } + bool IsDead() const { + RAY_CHECK(false) << "Method unused"; + return false; + } + void MarkBlocked() { RAY_CHECK(false) << "Method unused"; } + void MarkUnblocked() { RAY_CHECK(false) << "Method unused"; } + bool IsBlocked() const { + RAY_CHECK(false) << "Method unused"; + return false; + } + + Process GetProcess() const { + RAY_CHECK(false) << "Method unused"; + return Process::CreateNewDummy(); + } + void SetProcess(Process proc) { RAY_CHECK(false) << "Method unused"; } + Language GetLanguage() const { + RAY_CHECK(false) << "Method unused"; + return Language::PYTHON; + } + + void Connect(int port) { RAY_CHECK(false) << "Method unused"; } + + int AssignedPort() const { + RAY_CHECK(false) << "Method unused"; + return -1; + } + void SetAssignedPort(int port) { RAY_CHECK(false) << "Method unused"; } + const TaskID &GetAssignedTaskId() const { + RAY_CHECK(false) << "Method unused"; + return TaskID::Nil(); + } + bool AddBlockedTaskId(const TaskID &task_id) { + RAY_CHECK(false) << "Method unused"; + return false; + } + bool RemoveBlockedTaskId(const TaskID &task_id) { + RAY_CHECK(false) << "Method unused"; + return false; + } + const std::unordered_set &GetBlockedTaskIds() const { + RAY_CHECK(false) << "Method unused"; + auto *t = new std::unordered_set(); + return *t; + } + const JobID &GetAssignedJobId() const { + RAY_CHECK(false) << "Method unused"; + return JobID::Nil(); + } + void AssignActorId(const ActorID &actor_id) { RAY_CHECK(false) << "Method unused"; } + const ActorID &GetActorId() const { + RAY_CHECK(false) << "Method unused"; + return ActorID::Nil(); + } + void MarkDetachedActor() { RAY_CHECK(false) << "Method unused"; } + bool IsDetachedActor() const { + RAY_CHECK(false) << "Method unused"; + return false; + } + const std::shared_ptr Connection() const { + RAY_CHECK(false) << "Method unused"; + return nullptr; + } + const rpc::Address &GetOwnerAddress() const { + RAY_CHECK(false) << "Method unused"; + return address_; + } + + const ResourceIdSet &GetLifetimeResourceIds() const { + RAY_CHECK(false) << "Method unused"; + auto *t = new ResourceIdSet(); + return *t; + } + void SetLifetimeResourceIds(ResourceIdSet &resource_ids) { + RAY_CHECK(false) << "Method unused"; + } + void ResetLifetimeResourceIds() { RAY_CHECK(false) << "Method unused"; } + + const ResourceIdSet &GetTaskResourceIds() const { + RAY_CHECK(false) << "Method unused"; + auto *t = new ResourceIdSet(); + return *t; + } + void SetTaskResourceIds(ResourceIdSet &resource_ids) { + RAY_CHECK(false) << "Method unused"; + } + void ResetTaskResourceIds() { RAY_CHECK(false) << "Method unused"; } + ResourceIdSet ReleaseTaskCpuResources() { + RAY_CHECK(false) << "Method unused"; + auto *t = new ResourceIdSet(); + return *t; + } + void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) { + RAY_CHECK(false) << "Method unused"; + } + + Status AssignTask(const Task &task, const ResourceIdSet &resource_id_set) { + RAY_CHECK(false) << "Method unused"; + Status s; + return s; + } + void DirectActorCallArgWaitComplete(int64_t tag) { + RAY_CHECK(false) << "Method unused"; + } + + void ClearAllocatedInstances() { allocated_instances_ = nullptr; } + + void ClearLifetimeAllocatedInstances() { RAY_CHECK(false) << "Method unused"; } + + void SetBorrowedCPUInstances(std::vector &cpu_instances) { + borrowed_cpu_instances_ = cpu_instances; + } + + const PlacementGroupID &GetPlacementGroupId() const { + RAY_CHECK(false) << "Method unused"; + return PlacementGroupID::Nil(); + } + + void SetPlacementGroupId(const PlacementGroupID &placement_group_id) { + RAY_CHECK(false) << "Method unused"; + } + + std::vector &GetBorrowedCPUInstances() { return borrowed_cpu_instances_; } + + void ClearBorrowedCPUInstances() { RAY_CHECK(false) << "Method unused"; } + + Task &GetAssignedTask() { + RAY_CHECK(false) << "Method unused"; + auto *t = new Task(); + return *t; + } + + bool IsRegistered() { + RAY_CHECK(false) << "Method unused"; + return false; + } + + rpc::CoreWorkerClientInterface *rpc_client() { + RAY_CHECK(false) << "Method unused"; + return nullptr; + } + + private: + WorkerID worker_id_; + int port_; + rpc::Address address_; + std::shared_ptr allocated_instances_; + std::shared_ptr lifetime_allocated_instances_; + std::vector borrowed_cpu_instances_; +}; + +} // namespace raylet + +} // namespace ray diff --git a/src/ray/raylet/worker.h b/src/ray/raylet/worker.h index 856991e86..f739c98cb 100644 --- a/src/ray/raylet/worker.h +++ b/src/ray/raylet/worker.h @@ -111,7 +111,7 @@ class WorkerInterface { virtual bool IsRegistered() = 0; - virtual rpc::CoreWorkerClient *rpc_client() = 0; + virtual rpc::CoreWorkerClientInterface *rpc_client() = 0; }; /// Worker class encapsulates the implementation details of a worker. A worker @@ -211,7 +211,7 @@ class Worker : public WorkerInterface { bool IsRegistered() { return rpc_client_ != nullptr; } - rpc::CoreWorkerClient *rpc_client() { + rpc::CoreWorkerClientInterface *rpc_client() { RAY_CHECK(IsRegistered()); return rpc_client_.get(); } diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h index 5fbfce630..b325cf33a 100644 --- a/src/ray/raylet/worker_pool.h +++ b/src/ray/raylet/worker_pool.h @@ -57,6 +57,28 @@ class WorkerPoolInterface { virtual ~WorkerPoolInterface(){}; }; +/// \class IOWorkerPoolInterface +/// +/// Used for object spilling manager unit tests. +class IOWorkerPoolInterface { + public: + /// Add an idle I/O worker to the pool. + /// + /// \param worker The idle I/O worker to add. + virtual void PushIOWorker(const std::shared_ptr &worker) = 0; + + /// Pop an idle I/O worker from the pool and trigger a callback when + /// an I/O worker is available. + /// The caller is responsible for pushing the worker back onto the + /// pool once the worker has completed its work. + /// + /// \param callback The callback that returns an available I/O worker. + virtual void PopIOWorker( + std::function)> callback) = 0; + + virtual ~IOWorkerPoolInterface(){}; +}; + class WorkerInterface; class Worker; @@ -64,7 +86,7 @@ class Worker; /// /// The WorkerPool is responsible for managing a pool of Workers. Each Worker /// is a container for a unit of work. -class WorkerPool : public WorkerPoolInterface { +class WorkerPool : public WorkerPoolInterface, public IOWorkerPoolInterface { public: /// Create a pool and asynchronously start at least the specified number of workers per /// language.