From 407302f93af41dcdb20b5c7f5ae031073f6b1006 Mon Sep 17 00:00:00 2001 From: Clark Zinzow Date: Wed, 3 Feb 2021 15:16:42 -0700 Subject: [PATCH] [Core] Ownership-based Object Directory - Changed infinite short-poll location subscription to long-poll. (#13841) --- src/ray/core_worker/core_worker.cc | 26 ++++--- src/ray/core_worker/reference_count.cc | 44 ++++++++++- src/ray/core_worker/reference_count.h | 41 ++++++++-- .../ownership_based_object_directory.cc | 76 ++++++++++++++----- src/ray/protobuf/core_worker.proto | 6 ++ 5 files changed, 155 insertions(+), 38 deletions(-) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index a8c2e8557..6c8287c15 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -2219,19 +2219,25 @@ void CoreWorker::HandleGetObjectLocationsOwner( return; } auto object_id = ObjectID::FromBinary(request.object_id()); - absl::optional> node_ids = - reference_counter_->GetObjectLocations(object_id); - Status status; - if (node_ids.has_value()) { - for (const auto &node_id : node_ids.value()) { + const auto &callback = [object_id, reply, send_reply_callback]( + const absl::flat_hash_set &locations, + int64_t object_size, int64_t current_version) { + RAY_LOG(DEBUG) << "Replying to HandleGetObjectLocationsOwner for " << object_id + << " with location update version " << current_version << ", " + << locations.size() << " locations, and " << object_size + << " object size."; + for (const auto &node_id : locations) { reply->add_node_ids(node_id.Binary()); } - status = Status::OK(); - } else { - status = Status::ObjectNotFound("Object " + object_id.Hex() + " not found"); + reply->set_object_size(object_size); + reply->set_current_version(current_version); + send_reply_callback(Status::OK(), nullptr, nullptr); + }; + auto status = reference_counter_->SubscribeObjectLocations( + object_id, request.last_version(), callback); + if (!status.ok()) { + send_reply_callback(status, nullptr, nullptr); } - reply->set_object_size(reference_counter_->GetObjectSize(object_id)); - send_reply_callback(status, nullptr, nullptr); } void CoreWorker::HandleWaitForRefRemoved(const rpc::WaitForRefRemovedRequest &request, diff --git a/src/ray/core_worker/reference_count.cc b/src/ray/core_worker/reference_count.cc index ba2e20994..a38a98d80 100644 --- a/src/ray/core_worker/reference_count.cc +++ b/src/ray/core_worker/reference_count.cc @@ -185,6 +185,7 @@ void ReferenceCounter::UpdateObjectSize(const ObjectID &object_id, int64_t objec auto it = object_id_refs_.find(object_id); if (it != object_id_refs_.end()) { it->second.object_size = object_size; + PushToLocationSubscribers(it); } } @@ -915,11 +916,12 @@ bool ReferenceCounter::AddObjectLocation(const ObjectID &object_id, absl::MutexLock lock(&mutex_); auto it = object_id_refs_.find(object_id); if (it == object_id_refs_.end()) { - RAY_LOG(WARNING) << "Tried to add an object location for an object " << object_id - << " that doesn't exist in the reference table"; + RAY_LOG(INFO) << "Tried to add an object location for an object " << object_id + << " that doesn't exist in the reference table"; return false; } it->second.locations.insert(node_id); + PushToLocationSubscribers(it); return true; } @@ -928,11 +930,12 @@ bool ReferenceCounter::RemoveObjectLocation(const ObjectID &object_id, absl::MutexLock lock(&mutex_); auto it = object_id_refs_.find(object_id); if (it == object_id_refs_.end()) { - RAY_LOG(WARNING) << "Tried to remove an object location for an object " << object_id - << " that doesn't exist in the reference table"; + RAY_LOG(INFO) << "Tried to remove an object location for an object " << object_id + << " that doesn't exist in the reference table"; return false; } it->second.locations.erase(node_id); + PushToLocationSubscribers(it); return true; } @@ -1003,6 +1006,39 @@ absl::optional ReferenceCounter::GetLocalityData( return locality_data; } +void ReferenceCounter::PushToLocationSubscribers(ReferenceTable::iterator it) { + const auto callbacks = it->second.location_subscription_callbacks; + it->second.location_subscription_callbacks.clear(); + it->second.location_version++; + for (const auto callback : callbacks) { + callback(it->second.locations, it->second.object_size, it->second.location_version); + } +} + +Status ReferenceCounter::SubscribeObjectLocations( + const ObjectID &object_id, int64_t last_location_version, + const LocationSubscriptionCallback &callback) { + absl::MutexLock lock(&mutex_); + auto it = object_id_refs_.find(object_id); + if (it == object_id_refs_.end()) { + RAY_LOG(INFO) << "Tried to register a location subscriber for an object " << object_id + << " that doesn't exist in the reference table." + << " The object has probably already been freed."; + return Status::ObjectNotFound("Object " + object_id.Hex() + " not found"); + } + + if (last_location_version < it->second.location_version) { + // If the last location version is less than the current location version, we + // already have location data that the subscriber hasn't seen yet, so we immediately + // invoke the callback. + callback(it->second.locations, it->second.object_size, it->second.location_version); + } else { + // Otherwise, save the callback for later invocation. + it->second.location_subscription_callbacks.push_back(callback); + } + return Status::OK(); +} + ReferenceCounter::Reference ReferenceCounter::Reference::FromProto( const rpc::ObjectReferenceCount &ref_count) { Reference ref; diff --git a/src/ray/core_worker/reference_count.h b/src/ray/core_worker/reference_count.h index 9c0576393..014b94714 100644 --- a/src/ray/core_worker/reference_count.h +++ b/src/ray/core_worker/reference_count.h @@ -49,6 +49,10 @@ class ReferenceCounterInterface { virtual ~ReferenceCounterInterface() {} }; +// Callback for location subscriptions. +using LocationSubscriptionCallback = + std::function &, int64_t, int64_t)>; + /// Class used by the core worker to keep track of ObjectID reference counts for garbage /// collection. This class is thread safe. class ReferenceCounter : public ReferenceCounterInterface, @@ -397,6 +401,19 @@ class ReferenceCounter : public ReferenceCounterInterface, absl::optional> GetObjectLocations( const ObjectID &object_id) LOCKS_EXCLUDED(mutex_); + /// Subscribe to object location changes that are more recent than the given version. + /// The provided callback will be invoked when new locations become available. + /// + /// \param[in] object_id The object whose locations we want. + /// \param[in] last_location_version The version of the last location update the + /// caller received. Only more recent location updates will be returned. + /// \param[in] callback The callback to invoke with the location update. + /// \return The status of the location get. + Status SubscribeObjectLocations(const ObjectID &object_id, + int64_t last_location_version, + const LocationSubscriptionCallback &callback) + LOCKS_EXCLUDED(mutex_); + /// Get an object's size. This will return 0 if the object is out of scope. /// /// \param[in] object_id The object whose size to get. @@ -492,13 +509,17 @@ class ReferenceCounter : public ReferenceCounterInterface, /// process is a borrower, the borrower must add the owner's address before /// using the ObjectID. absl::optional owner_address; - // If this object is owned by us and stored in plasma, and reference - // counting is enabled, then some raylet must be pinning the object value. - // This is the address of that raylet. + /// If this object is owned by us and stored in plasma, and reference + /// counting is enabled, then some raylet must be pinning the object value. + /// This is the address of that raylet. absl::optional pinned_at_raylet_id; - // If this object is owned by us and stored in plasma, this contains all - // object locations. + /// If this object is owned by us and stored in plasma, this contains all + /// object locations. absl::flat_hash_set locations; + /// A logical counter for object location updates, used for object location + /// subscriptions. Subscribers use -1 to indicate that they want us to + /// immediately send them the current location data. + int64_t location_version = 0; // Whether this object can be reconstructed via lineage. If false, then the // object's value will be pinned as long as it is referenced by any other // object's lineage. @@ -565,7 +586,9 @@ class ReferenceCounter : public ReferenceCounterInterface, size_t lineage_ref_count = 0; /// Whether this object has been spilled to external storage. bool spilled = false; - + /// Location subscription callbacks registered by async location get requests. + /// These will be invoked whenever locations or object_size are changed. + std::vector location_subscription_callbacks; /// Callback that will be called when this ObjectID no longer has /// references. std::function on_delete; @@ -689,6 +712,12 @@ class ReferenceCounter : public ReferenceCounterInterface, void ReleaseLineageReferencesInternal(const std::vector &argument_ids) EXCLUSIVE_LOCKS_REQUIRED(mutex_); + /// Pushes location updates to subscribers of a particular reference, invoking all + /// callbacks registered for the reference by GetLocationsAsync calls. This method + /// also increments the reference's location version counter. + void PushToLocationSubscribers(ReferenceTable::iterator it) + EXCLUSIVE_LOCKS_REQUIRED(mutex_); + /// Address of our RPC server. This is used to determine whether we own a /// given object or not, by comparing our WorkerID with the WorkerID of the /// object's owner. diff --git a/src/ray/object_manager/ownership_based_object_directory.cc b/src/ray/object_manager/ownership_based_object_directory.cc index a17d3dfc6..3f2ccc540 100644 --- a/src/ray/object_manager/ownership_based_object_directory.cc +++ b/src/ray/object_manager/ownership_based_object_directory.cc @@ -80,11 +80,18 @@ ray::Status OwnershipBasedObjectDirectory::ReportObjectAdded( request.set_node_id(node_id.Binary()); rpc_client->AddObjectLocationOwner( - request, [worker_id, object_id](Status status, - const rpc::AddObjectLocationOwnerReply &reply) { + request, [worker_id, object_id, node_id]( + Status status, const rpc::AddObjectLocationOwnerReply &reply) { if (!status.ok()) { - RAY_LOG(ERROR) << "Worker " << worker_id << " failed to add the location for " - << object_id; + if (status.IsObjectNotFound()) { + RAY_LOG(INFO) << "Worker " << worker_id << " failed to add the location " + << node_id << " for " << object_id + << " because the owner no longer has the object; we assume the " + "object was evicted."; + } else { + RAY_LOG(INFO) << "Worker " << worker_id << " failed to add the location " + << node_id << " for " << object_id << ": " << status.ToString(); + } } }); return Status::OK(); @@ -108,11 +115,18 @@ ray::Status OwnershipBasedObjectDirectory::ReportObjectRemoved( request.set_node_id(node_id.Binary()); rpc_client->RemoveObjectLocationOwner( - request, [worker_id, object_id](Status status, - const rpc::RemoveObjectLocationOwnerReply &reply) { + request, [worker_id, object_id, node_id]( + Status status, const rpc::RemoveObjectLocationOwnerReply &reply) { if (!status.ok()) { - RAY_LOG(ERROR) << "Worker " << worker_id - << " failed to remove the location for " << object_id; + if (status.IsObjectNotFound()) { + RAY_LOG(INFO) << "Worker " << worker_id << " failed to remove the location " + << node_id << " for " << object_id + << " because the owner no longer has the object; we assume the " + "object was freed."; + } else { + RAY_LOG(INFO) << "Worker " << worker_id << " failed to remove the location " + << node_id << " for " << object_id << ": " << status.ToString(); + } } }); return Status::OK(); @@ -121,22 +135,36 @@ ray::Status OwnershipBasedObjectDirectory::ReportObjectRemoved( void OwnershipBasedObjectDirectory::SubscriptionCallback( ObjectID object_id, WorkerID worker_id, Status status, const rpc::GetObjectLocationsOwnerReply &reply) { + // Objects are added to this map in SubscribeObjectLocations. auto it = listeners_.find(object_id); + // Do nothing for objects we are not listening for. if (it == listeners_.end()) { return; } - - if (reply.object_size() > 0) { - it->second.object_size = reply.object_size(); - } - std::unordered_set node_ids; - for (auto const &node_id : reply.node_ids()) { - node_ids.emplace(NodeID::FromBinary(node_id)); + + // Once this flag is set to true, it should never go back to false. + it->second.subscribed = true; + + if (!status.ok()) { + RAY_LOG(INFO) << "Worker " << worker_id << " failed to return location updates to " + << "subscribers for " << object_id << ": " << status.ToString() + << ", assuming that the object was freed or evicted."; + it->second.object_size = 0; + } else { + if (reply.object_size() > 0) { + it->second.object_size = reply.object_size(); + } + + for (auto const &node_id : reply.node_ids()) { + node_ids.emplace(NodeID::FromBinary(node_id)); + } + FilterRemovedNodes(gcs_client_, &node_ids); } - FilterRemovedNodes(gcs_client_, &node_ids); - if (node_ids != it->second.current_object_locations) { + if (node_ids != it->second.current_object_locations || !status.ok()) { it->second.current_object_locations = std::move(node_ids); + // Copy the callbacks so that the callbacks can unsubscribe without interrupting + // looping over the callbacks. auto callbacks = it->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 @@ -154,7 +182,7 @@ void OwnershipBasedObjectDirectory::SubscriptionCallback( rpc::GetObjectLocationsOwnerRequest request; request.set_intended_worker_id(worker_id.Binary()); request.set_object_id(object_id.Binary()); - // TODO(zhuohan): Fix this infinite loop. + request.set_last_version(reply.current_version()); worker_it->second->GetObjectLocationsOwner( request, std::bind(&OwnershipBasedObjectDirectory::SubscriptionCallback, this, object_id, @@ -176,6 +204,7 @@ ray::Status OwnershipBasedObjectDirectory::SubscribeObjectLocations( rpc::GetObjectLocationsOwnerRequest request; request.set_intended_worker_id(owner_address.worker_id()); request.set_object_id(object_id.Binary()); + request.set_last_version(-1); rpc_client->GetObjectLocationsOwner( request, std::bind(&OwnershipBasedObjectDirectory::SubscriptionCallback, this, object_id, @@ -188,6 +217,16 @@ ray::Status OwnershipBasedObjectDirectory::SubscribeObjectLocations( return Status::OK(); } listener_state.callbacks.emplace(callback_id, callback); + + // If we previously received some notifications about the object's locations, + // immediately notify the caller of the current known locations. + if (listener_state.subscribed) { + auto &locations = listener_state.current_object_locations; + auto object_size = it->second.object_size; + io_service_.post([callback, locations, object_size, object_id]() { + callback(object_id, locations, "", NodeID::Nil(), object_size); + }); + } return Status::OK(); } @@ -221,6 +260,7 @@ ray::Status OwnershipBasedObjectDirectory::LookupLocations( rpc::GetObjectLocationsOwnerRequest request; request.set_intended_worker_id(owner_address.worker_id()); request.set_object_id(object_id.Binary()); + request.set_last_version(-1); rpc_client->GetObjectLocationsOwner( request, [this, worker_id, object_id, callback]( diff --git a/src/ray/protobuf/core_worker.proto b/src/ray/protobuf/core_worker.proto index 43a3a6674..ef5f97302 100644 --- a/src/ray/protobuf/core_worker.proto +++ b/src/ray/protobuf/core_worker.proto @@ -182,11 +182,17 @@ message RemoveObjectLocationOwnerReply { message GetObjectLocationsOwnerRequest { bytes intended_worker_id = 1; bytes object_id = 2; + // The version of the last location update. Only updates more recent than this version + // will be returned. -1 indicates that the current location data should + // always be returned. + int64 last_version = 3; } message GetObjectLocationsOwnerReply { repeated bytes node_ids = 1; uint64 object_size = 2; + // The version of the returned location updates. + int64 current_version = 3; } message KillActorRequest {