From a09997dc9e180628a1e521d76a940b3654a6c991 Mon Sep 17 00:00:00 2001 From: "Siyuan (Ryans) Zhuang" Date: Wed, 20 Jan 2021 12:01:44 -0800 Subject: [PATCH] [Core] Remove 'PlasmaBuffer' in the buffer header (#13188) --- src/ray/common/buffer.h | 32 +-------- .../store_provider/plasma_store_provider.cc | 72 ++++++++++--------- .../store_provider/plasma_store_provider.h | 66 +++++++++++++---- src/ray/core_worker/test/core_worker_test.cc | 2 +- src/ray/raylet/node_manager.cc | 5 +- 5 files changed, 94 insertions(+), 83 deletions(-) diff --git a/src/ray/common/buffer.h b/src/ray/common/buffer.h index 2bcc92751..6dc25033a 100644 --- a/src/ray/common/buffer.h +++ b/src/ray/common/buffer.h @@ -152,7 +152,7 @@ class SharedMemoryBuffer : public Buffer { size_t Size() const override { return size_; } - bool OwnsData() const override { return false; } + bool OwnsData() const override { return true; } bool IsPlasmaBuffer() const override { return true; } @@ -172,34 +172,4 @@ class SharedMemoryBuffer : public Buffer { std::shared_ptr parent_; }; -/// Represents a byte buffer for plasma object. This can be used to hold the -/// reference to a plasma object (via the underlying plasma::PlasmaBuffer). -class PlasmaBuffer : public Buffer { - public: - PlasmaBuffer(std::shared_ptr buffer, - std::function on_delete = nullptr) - : buffer_(buffer), on_delete_(on_delete) {} - - uint8_t *Data() const override { return buffer_->Data(); } - - size_t Size() const override { return buffer_->Size(); } - - bool OwnsData() const override { return true; } - - bool IsPlasmaBuffer() const override { return true; } - - ~PlasmaBuffer() { - if (on_delete_ != nullptr) { - on_delete_(this); - } - }; - - private: - /// shared_ptr to a buffer which can potentially hold a reference - /// for the object (when it's a plasma::PlasmaBuffer). - std::shared_ptr buffer_; - /// Callback to run on destruction. - std::function on_delete_; -}; - } // namespace ray diff --git a/src/ray/core_worker/store_provider/plasma_store_provider.cc b/src/ray/core_worker/store_provider/plasma_store_provider.cc index 0d6789dc4..831f2629a 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.cc +++ b/src/ray/core_worker/store_provider/plasma_store_provider.cc @@ -21,6 +21,36 @@ namespace ray { +void BufferTracker::Record(const ObjectID &object_id, TrackedBuffer *buffer, + const std::string &call_site) { + absl::MutexLock lock(&active_buffers_mutex_); + active_buffers_[std::make_pair(object_id, buffer)] = call_site; +} + +void BufferTracker::Release(const ObjectID &object_id, TrackedBuffer *buffer) { + absl::MutexLock lock(&active_buffers_mutex_); + auto key = std::make_pair(object_id, buffer); + RAY_CHECK(active_buffers_.contains(key)); + active_buffers_.erase(key); +} + +absl::flat_hash_map> +BufferTracker::UsedObjects() const { + absl::flat_hash_map> used; + absl::MutexLock lock(&active_buffers_mutex_); + for (const auto &entry : active_buffers_) { + auto it = used.find(entry.first.first); + if (it != used.end()) { + // Prefer to keep entries that have non-empty callsites. + if (!it->second.second.empty()) { + continue; + } + } + used[entry.first.first] = std::make_pair(entry.first.second->Size(), entry.second); + } + return used; +} + CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider( const std::string &store_socket, const std::shared_ptr raylet_client, @@ -85,13 +115,12 @@ Status CoreWorkerPlasmaStoreProvider::Create(const std::shared_ptr &meta const rpc::Address &owner_address, std::shared_ptr *data) { Status status; - std::shared_ptr plasma_buffer; uint64_t retry_with_request_id = 0; { std::lock_guard guard(store_client_mutex_); status = store_client_.Create( object_id, owner_address, data_size, metadata ? metadata->Data() : nullptr, - metadata ? metadata->Size() : 0, &retry_with_request_id, &plasma_buffer, + metadata ? metadata->Size() : 0, &retry_with_request_id, data, /*device_num=*/0); } @@ -104,7 +133,7 @@ Status CoreWorkerPlasmaStoreProvider::Create(const std::shared_ptr &meta << retry_with_request_id; status = store_client_.RetryCreate(object_id, retry_with_request_id, metadata ? metadata->Data() : nullptr, - &retry_with_request_id, &plasma_buffer); + &retry_with_request_id, data); } } @@ -129,7 +158,6 @@ Status CoreWorkerPlasmaStoreProvider::Create(const std::shared_ptr &meta status = Status::OK(); } else { RAY_RETURN_NOT_OK(status); - *data = std::make_shared(PlasmaBuffer(plasma_buffer)); } return status; } @@ -171,27 +199,17 @@ Status CoreWorkerPlasmaStoreProvider::FetchAndGetFromPlasmaStore( for (size_t i = 0; i < plasma_results.size(); i++) { if (plasma_results[i].data != nullptr || plasma_results[i].metadata != nullptr) { const auto &object_id = batch_ids[i]; - std::shared_ptr data = nullptr; - std::shared_ptr metadata = nullptr; + std::shared_ptr data = nullptr; + std::shared_ptr metadata = nullptr; if (plasma_results[i].data && plasma_results[i].data->Size()) { // We track the set of active data buffers in active_buffers_. On destruction, // the buffer entry will be removed from the set via callback. - std::shared_ptr tracker = buffer_tracker_; - data = std::make_shared( - plasma_results[i].data, [tracker, object_id](PlasmaBuffer *this_buffer) { - absl::MutexLock lock(&tracker->active_buffers_mutex_); - auto key = std::make_pair(object_id, this_buffer); - RAY_CHECK(tracker->active_buffers_.contains(key)); - tracker->active_buffers_.erase(key); - }); - auto call_site = get_current_call_site_(); - { - absl::MutexLock lock(&tracker->active_buffers_mutex_); - tracker->active_buffers_[std::make_pair(object_id, data.get())] = call_site; - } + data = std::make_shared(plasma_results[i].data, buffer_tracker_, + object_id); + buffer_tracker_->Record(object_id, data.get(), get_current_call_site_()); } if (plasma_results[i].metadata && plasma_results[i].metadata->Size()) { - metadata = std::make_shared(plasma_results[i].metadata); + metadata = plasma_results[i].metadata; } const auto result_object = std::make_shared(data, metadata, std::vector()); @@ -373,19 +391,7 @@ std::string CoreWorkerPlasmaStoreProvider::MemoryUsageString() { absl::flat_hash_map> CoreWorkerPlasmaStoreProvider::UsedObjectsList() const { - absl::flat_hash_map> used; - absl::MutexLock lock(&buffer_tracker_->active_buffers_mutex_); - for (const auto &entry : buffer_tracker_->active_buffers_) { - auto it = used.find(entry.first.first); - if (it != used.end()) { - // Prefer to keep entries that have non-empty callsites. - if (!it->second.second.empty()) { - continue; - } - } - used[entry.first.first] = std::make_pair(entry.first.second->Size(), entry.second); - } - return used; + return buffer_tracker_->UsedObjects(); } void CoreWorkerPlasmaStoreProvider::WarnIfAttemptedTooManyTimes( diff --git a/src/ray/core_worker/store_provider/plasma_store_provider.h b/src/ray/core_worker/store_provider/plasma_store_provider.h index e9c7a23ee..2282a09a9 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.h +++ b/src/ray/core_worker/store_provider/plasma_store_provider.h @@ -27,6 +27,57 @@ namespace ray { +class TrackedBuffer; + +// Active buffers tracker. This must be allocated as a separate structure since its +// lifetime can exceed that of the store provider due to TrackedBuffer. +class BufferTracker { + public: + // Track an object. + void Record(const ObjectID &object_id, TrackedBuffer *buffer, + const std::string &call_site); + // Release an object from tracking. + void Release(const ObjectID &object_id, TrackedBuffer *buffer); + // List tracked objects. + absl::flat_hash_map> UsedObjects() const; + + private: + // Guards the active buffers map. This mutex may be acquired during TrackedBuffer + // destruction. + mutable absl::Mutex active_buffers_mutex_; + // Mapping of live object buffers to their creation call site. Destroyed buffers are + // automatically removed from this list via destructor. The map key uniquely + // identifies a buffer. It should not be a shared ptr since that would keep the Buffer + // alive forever (i.e., this is a weak ref map). + absl::flat_hash_map, std::string> active_buffers_ + GUARDED_BY(active_buffers_mutex_); +}; + +/// This can be used to hold the reference to a buffer. +class TrackedBuffer : public Buffer { + public: + TrackedBuffer(std::shared_ptr buffer, + const std::shared_ptr &tracker, const ObjectID &object_id) + : buffer_(buffer), tracker_(tracker), object_id_(object_id) {} + + uint8_t *Data() const override { return buffer_->Data(); } + + size_t Size() const override { return buffer_->Size(); } + + bool OwnsData() const override { return true; } + + bool IsPlasmaBuffer() const override { return true; } + + ~TrackedBuffer() { tracker_->Release(object_id_, this); } + + private: + /// shared_ptr to a buffer which can potentially hold a reference + /// for the object (when it's a SharedMemoryBuffer). + std::shared_ptr buffer_; + std::shared_ptr tracker_; + ObjectID object_id_; +}; + /// The class provides implementations for accessing plasma store, which includes both /// local and remote stores. Local access goes is done via a /// CoreWorkerLocalPlasmaStoreProvider and remote access goes through the raylet. @@ -154,21 +205,6 @@ class CoreWorkerPlasmaStoreProvider { std::function check_signals_; std::function get_current_call_site_; uint32_t object_store_full_delay_ms_; - - // Active buffers tracker. This must be allocated as a separate structure since its - // lifetime can exceed that of the store provider due to callback references. - struct BufferTracker { - // Guards the active buffers map. This mutex may be acquired during PlasmaBuffer - // destruction. - mutable absl::Mutex active_buffers_mutex_; - // Mapping of live object buffers to their creation call site. Destroyed buffers are - // automatically removed from this list via destructor callback. The map key uniquely - // identifies a buffer. It should not be a shared ptr since that would keep the Buffer - // alive forever (i.e., this is a weak ref map). - absl::flat_hash_map, std::string> active_buffers_ - GUARDED_BY(active_buffers_mutex_); - }; - // Pointer to the shared buffer tracker. std::shared_ptr buffer_tracker_; }; diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc index 0c4d69149..82ea82617 100644 --- a/src/ray/core_worker/test/core_worker_test.cc +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -820,7 +820,7 @@ TEST_F(SingleNodeTest, TestObjectInterface) { ASSERT_EQ(wait_results, std::vector({true, true, false})); // Test Delete(). - // clear the reference held by PlasmaBuffer. + // clear the reference held by TrackedBuffer. results.clear(); RAY_CHECK_OK(core_worker.Delete(ids, true)); diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 7d4f97d0e..1b8c50c58 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -2368,9 +2368,8 @@ void NodeManager::HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request, if (plasma_results[i].data == nullptr) { 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), {}))); + objects.emplace_back(std::unique_ptr( + new RayObject(plasma_results[i].data, plasma_results[i].metadata, {}))); } } local_object_manager_.PinObjects(object_ids, std::move(objects));