diff --git a/src/ray/core_worker/common.h b/src/ray/core_worker/common.h index aeb433065..631ada2ba 100644 --- a/src/ray/core_worker/common.h +++ b/src/ray/core_worker/common.h @@ -64,7 +64,7 @@ class TaskArg { const std::shared_ptr data_; }; -enum class StoreProviderType { LOCAL_PLASMA, PLASMA, MEMORY }; +enum class StoreProviderType { PLASMA, MEMORY }; enum class TaskTransportType { RAYLET, DIRECT_ACTOR }; diff --git a/src/ray/core_worker/object_interface.cc b/src/ray/core_worker/object_interface.cc index cb4e68a41..406c5edfc 100644 --- a/src/ray/core_worker/object_interface.cc +++ b/src/ray/core_worker/object_interface.cc @@ -2,7 +2,6 @@ #include "ray/common/ray_config.h" #include "ray/core_worker/object_interface.h" -#include "ray/core_worker/store_provider/local_plasma_provider.h" #include "ray/core_worker/store_provider/memory_store_provider.h" #include "ray/core_worker/store_provider/plasma_store_provider.h" @@ -43,7 +42,6 @@ CoreWorkerObjectInterface::CoreWorkerObjectInterface( raylet_client_(raylet_client), store_socket_(store_socket), memory_store_(std::make_shared()) { - AddStoreProvider(StoreProviderType::LOCAL_PLASMA); AddStoreProvider(StoreProviderType::PLASMA); AddStoreProvider(StoreProviderType::MEMORY); } @@ -76,17 +74,17 @@ Status CoreWorkerObjectInterface::Get(const std::vector &ids, object_ids_per_store_provider; GroupObjectIdsByStoreProvider(ids, &object_ids_per_store_provider); - std::unordered_map> objects; + std::unordered_map> result_map; auto remaining_timeout_ms = timeout_ms; // Re-order the list so that we always get from plasma store provider first, // since it uses a loop of `FetchOrReconstruct` and plasma `Get`, it's not // desirable if other store providers use up the timeout and leaves no time // for plasma provider to reconstruct the objects as necessary. - std::list< - std::pair>>> + std::list>>> ids_per_provider; - for (auto &entry : object_ids_per_store_provider) { + for (const auto &entry : object_ids_per_store_provider) { auto list_entry = std::make_pair(entry.first, std::ref(entry.second)); if (entry.first == StoreProviderType::PLASMA) { ids_per_provider.emplace_front(list_entry); @@ -97,10 +95,11 @@ Status CoreWorkerObjectInterface::Get(const std::vector &ids, // Note that if one store provider uses up the timeout, we will still try the others // with a timeout of 0. - for (const auto &entry : object_ids_per_store_provider) { + for (const auto &entry : ids_per_provider) { auto start_time = current_time_ms(); - RAY_RETURN_NOT_OK( - GetFromStoreProvider(entry.first, entry.second, remaining_timeout_ms, &objects)); + RAY_RETURN_NOT_OK(store_providers_[entry.first]->Get( + entry.second, remaining_timeout_ms, worker_context_.GetCurrentTaskID(), + &result_map)); if (remaining_timeout_ms > 0) { int64_t duration = current_time_ms() - start_time; remaining_timeout_ms = @@ -113,24 +112,8 @@ Status CoreWorkerObjectInterface::Get(const std::vector &ids, // they are in `ids`. When there are duplicate object ids, all the entries // for the same id are filled in. for (size_t i = 0; i < ids.size(); i++) { - (*results)[i] = objects[ids[i]]; - } - - return Status::OK(); -} - -Status CoreWorkerObjectInterface::GetFromStoreProvider( - StoreProviderType type, const std::unordered_set &object_ids, - int64_t timeout_ms, - std::unordered_map> *results) { - std::vector ids(object_ids.begin(), object_ids.end()); - if (!ids.empty()) { - std::vector> objects; - RAY_RETURN_NOT_OK(store_providers_[type]->Get( - ids, timeout_ms, worker_context_.GetCurrentTaskID(), &objects)); - RAY_CHECK(ids.size() == objects.size()); - for (size_t i = 0; i < objects.size(); i++) { - (*results).emplace(ids[i], objects[i]); + if (result_map.find(ids[i]) != result_map.end()) { + (*results)[i] = result_map[ids[i]]; } } @@ -256,9 +239,6 @@ void CoreWorkerObjectInterface::AddStoreProvider(StoreProviderType type) { std::unique_ptr CoreWorkerObjectInterface::CreateStoreProvider( StoreProviderType type) const { switch (type) { - case StoreProviderType::LOCAL_PLASMA: - return std::unique_ptr( - new CoreWorkerLocalPlasmaStoreProvider(store_socket_)); case StoreProviderType::PLASMA: return std::unique_ptr( new CoreWorkerPlasmaStoreProvider(store_socket_, raylet_client_)); diff --git a/src/ray/core_worker/object_interface.h b/src/ray/core_worker/object_interface.h index dab30a3b2..bcf23b30b 100644 --- a/src/ray/core_worker/object_interface.h +++ b/src/ray/core_worker/object_interface.h @@ -86,18 +86,6 @@ class CoreWorkerObjectInterface { &ids_per_provider, int64_t timeout_ms, int *num_objects, std::vector *results); - /// Helper function to get a list of objects from a specific store provider. - /// - /// \param[in] type The type of store provider to use. - /// \param[in] object_ids IDs of the objects to get. - /// \param[in] timeout_ms Timeout in milliseconds, wait infinitely if it's -1. - /// \param[out] results Result list of objects data. - /// \return Status. - Status GetFromStoreProvider( - StoreProviderType type, const std::unordered_set &object_ids, - int64_t timeout_ms, - std::unordered_map> *results); - /// Helper function to wait a list of objects from a specific store provider. /// /// \param[in] type The type of store provider to use. diff --git a/src/ray/core_worker/store_provider/local_plasma_provider.cc b/src/ray/core_worker/store_provider/local_plasma_provider.cc deleted file mode 100644 index 8295a9fd5..000000000 --- a/src/ray/core_worker/store_provider/local_plasma_provider.cc +++ /dev/null @@ -1,106 +0,0 @@ -#include "ray/core_worker/store_provider/local_plasma_provider.h" -#include "ray/common/ray_config.h" -#include "ray/core_worker/context.h" -#include "ray/core_worker/core_worker.h" -#include "ray/core_worker/object_interface.h" - -namespace ray { - -CoreWorkerLocalPlasmaStoreProvider::CoreWorkerLocalPlasmaStoreProvider( - const std::string &store_socket) { - RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket)); -} - -Status CoreWorkerLocalPlasmaStoreProvider::Put(const RayObject &object, - const ObjectID &object_id) { - auto plasma_id = object_id.ToPlasmaId(); - auto data = object.GetData(); - auto metadata = object.GetMetadata(); - std::shared_ptr out_buffer; - { - std::unique_lock guard(store_client_mutex_); - arrow::Status status = store_client_.Create( - plasma_id, data ? data->Size() : 0, metadata ? metadata->Data() : nullptr, - metadata ? metadata->Size() : 0, &out_buffer); - if (plasma::IsPlasmaObjectExists(status)) { - // TODO(hchen): Should we propagate this error out of `ObjectInterface::put`? - RAY_LOG(WARNING) << "Trying to put an object that already existed in plasma: " - << object_id << "."; - return Status::OK(); - } - RAY_ARROW_RETURN_NOT_OK(status); - } - - if (data != nullptr) { - memcpy(out_buffer->mutable_data(), data->Data(), data->Size()); - } - - { - std::unique_lock guard(store_client_mutex_); - RAY_ARROW_RETURN_NOT_OK(store_client_.Seal(plasma_id)); - RAY_ARROW_RETURN_NOT_OK(store_client_.Release(plasma_id)); - } - return Status::OK(); -} - -Status CoreWorkerLocalPlasmaStoreProvider::Get( - const std::vector &object_ids, int64_t timeout_ms, const TaskID &task_id, - std::vector> *results) { - std::vector plasma_ids; - plasma_ids.reserve(object_ids.size()); - for (const auto &object_id : object_ids) { - plasma_ids.push_back(object_id.ToPlasmaId()); - } - - std::vector object_buffers; - { - std::unique_lock guard(store_client_mutex_); - RAY_ARROW_RETURN_NOT_OK(store_client_.Get(plasma_ids, timeout_ms, &object_buffers)); - } - - (*results).resize(object_ids.size(), nullptr); - for (size_t i = 0; i < object_buffers.size(); i++) { - if (object_buffers[i].data != nullptr || object_buffers[i].metadata != nullptr) { - (*results)[i] = std::make_shared( - std::make_shared(object_buffers[i].data), - std::make_shared(object_buffers[i].metadata)); - } - } - - return Status::OK(); -} - -Status CoreWorkerLocalPlasmaStoreProvider::Wait(const std::vector &object_ids, - int num_objects, int64_t timeout_ms, - const TaskID &task_id, - std::vector *results) { - if (num_objects != static_cast(object_ids.size())) { - return Status::Invalid("num_objects should equal to number of items in object_ids"); - } - - std::vector> objects; - RAY_RETURN_NOT_OK(Get(object_ids, timeout_ms, task_id, &objects)); - - (*results).resize(object_ids.size()); - for (size_t i = 0; i < object_ids.size(); i++) { - (*results)[i] = (objects[i] != nullptr); - } - - return Status::OK(); -} - -Status CoreWorkerLocalPlasmaStoreProvider::Delete(const std::vector &object_ids, - bool local_only, - bool delete_creating_tasks) { - std::vector plasma_ids; - plasma_ids.reserve(object_ids.size()); - for (const auto &object_id : object_ids) { - plasma_ids.push_back(object_id.ToPlasmaId()); - } - - std::unique_lock guard(store_client_mutex_); - RAY_ARROW_RETURN_NOT_OK(store_client_.Delete(plasma_ids)); - return Status::OK(); -} - -} // namespace ray diff --git a/src/ray/core_worker/store_provider/local_plasma_provider.h b/src/ray/core_worker/store_provider/local_plasma_provider.h deleted file mode 100644 index d67b916b9..000000000 --- a/src/ray/core_worker/store_provider/local_plasma_provider.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef RAY_CORE_WORKER_LOCAL_PLASMA_STORE_PROVIDER_H -#define RAY_CORE_WORKER_LOCAL_PLASMA_STORE_PROVIDER_H - -#include "plasma/client.h" -#include "ray/common/buffer.h" -#include "ray/common/id.h" -#include "ray/common/status.h" -#include "ray/core_worker/common.h" -#include "ray/core_worker/store_provider/store_provider.h" -#include "ray/raylet/raylet_client.h" - -namespace ray { - -class CoreWorker; - -/// The class provides implementations for accessing local plasma store. -class CoreWorkerLocalPlasmaStoreProvider : public CoreWorkerStoreProvider { - public: - CoreWorkerLocalPlasmaStoreProvider(const std::string &store_socket); - - /// See `CoreWorkerStoreProvider::Put` for semantics. - Status Put(const RayObject &object, const ObjectID &object_id) override; - - /// See `CoreWorkerStoreProvider::Get` for semantics. - Status Get(const std::vector &ids, int64_t timeout_ms, const TaskID &task_id, - std::vector> *results) override; - - /// See `CoreWorkerStoreProvider::Wait` for semantics. - /// Note that `num_objects` must equal to number of items in `object_ids`. - Status Wait(const std::vector &object_ids, int num_objects, - int64_t timeout_ms, const TaskID &task_id, - std::vector *results) override; - - /// See `CoreWorkerStoreProvider::Delete` for semantics. - /// Note that `local_only` msut be true, and `delete_creating_tasks` must be false here. - Status Delete(const std::vector &object_ids, bool local_only = true, - bool delete_creating_tasks = false) override; - - private: - /// Plasma store client. - plasma::PlasmaClient store_client_; - - /// Mutex to protect store_client_. - std::mutex store_client_mutex_; -}; - -} // namespace ray - -#endif // RAY_CORE_WORKER_LOCAL_PLASMA_STORE_PROVIDER_H diff --git a/src/ray/core_worker/store_provider/memory_store_provider.cc b/src/ray/core_worker/store_provider/memory_store_provider.cc index e83a5a307..fb796a9b3 100644 --- a/src/ray/core_worker/store_provider/memory_store_provider.cc +++ b/src/ray/core_worker/store_provider/memory_store_provider.cc @@ -22,9 +22,20 @@ Status CoreWorkerMemoryStoreProvider::Put(const RayObject &object, } Status CoreWorkerMemoryStoreProvider::Get( - const std::vector &object_ids, int64_t timeout_ms, const TaskID &task_id, - std::vector> *results) { - return store_->Get(object_ids, object_ids.size(), timeout_ms, true, results); + const std::unordered_set &object_ids, int64_t timeout_ms, + const TaskID &task_id, + std::unordered_map> *results) { + const std::vector id_vector(object_ids.begin(), object_ids.end()); + std::vector> result_objects; + RAY_RETURN_NOT_OK( + store_->Get(id_vector, id_vector.size(), timeout_ms, true, &result_objects)); + + for (size_t i = 0; i < id_vector.size(); i++) { + if (result_objects[i] != nullptr) { + (*results)[id_vector[i]] = result_objects[i]; + } + } + return Status::OK(); } Status CoreWorkerMemoryStoreProvider::Wait(const std::vector &object_ids, @@ -34,15 +45,14 @@ Status CoreWorkerMemoryStoreProvider::Wait(const std::vector &object_i (*results).resize(object_ids.size(), false); std::vector> result_objects; - auto status = store_->Get(object_ids, num_objects, timeout_ms, false, &result_objects); - if (status.ok()) { - RAY_CHECK(result_objects.size() == object_ids.size()); - for (size_t i = 0; i < object_ids.size(); i++) { - (*results)[i] = (result_objects[i] != nullptr); - } + RAY_RETURN_NOT_OK( + store_->Get(object_ids, num_objects, timeout_ms, false, &result_objects)); + RAY_CHECK(result_objects.size() == object_ids.size()); + for (size_t i = 0; i < object_ids.size(); i++) { + (*results)[i] = (result_objects[i] != nullptr); } - return status; + return Status::OK(); } Status CoreWorkerMemoryStoreProvider::Delete(const std::vector &object_ids, diff --git a/src/ray/core_worker/store_provider/memory_store_provider.h b/src/ray/core_worker/store_provider/memory_store_provider.h index 4d072c832..066d7cf41 100644 --- a/src/ray/core_worker/store_provider/memory_store_provider.h +++ b/src/ray/core_worker/store_provider/memory_store_provider.h @@ -23,8 +23,9 @@ class CoreWorkerMemoryStoreProvider : public CoreWorkerStoreProvider { Status Put(const RayObject &object, const ObjectID &object_id) override; /// See `CoreWorkerStoreProvider::Get` for semantics. - Status Get(const std::vector &ids, int64_t timeout_ms, const TaskID &task_id, - std::vector> *results) override; + Status Get(const std::unordered_set &object_ids, int64_t timeout_ms, + const TaskID &task_id, + std::unordered_map> *results) override; /// See `CoreWorkerStoreProvider::Wait` for semantics. /// Note that `num_objects` must equal to number of items in `object_ids`. 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 39cf58b8d..54b850405 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.cc +++ b/src/ray/core_worker/store_provider/plasma_store_provider.cc @@ -9,47 +9,104 @@ namespace ray { CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider( const std::string &store_socket, std::unique_ptr &raylet_client) - : local_store_provider_(store_socket), raylet_client_(raylet_client) {} + : raylet_client_(raylet_client) { + RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket)); +} Status CoreWorkerPlasmaStoreProvider::Put(const RayObject &object, const ObjectID &object_id) { - return local_store_provider_.Put(object, object_id); + auto plasma_id = object_id.ToPlasmaId(); + auto data = object.GetData(); + auto metadata = object.GetMetadata(); + std::shared_ptr out_buffer; + { + std::unique_lock guard(store_client_mutex_); + arrow::Status status = store_client_.Create( + plasma_id, data ? data->Size() : 0, metadata ? metadata->Data() : nullptr, + metadata ? metadata->Size() : 0, &out_buffer); + if (plasma::IsPlasmaObjectExists(status)) { + // TODO(hchen): Should we propagate this error out of `ObjectInterface::put`? + RAY_LOG(WARNING) << "Trying to put an object that already existed in plasma: " + << object_id << "."; + return Status::OK(); + } + RAY_ARROW_RETURN_NOT_OK(status); + } + + if (data != nullptr) { + memcpy(out_buffer->mutable_data(), data->Data(), data->Size()); + } + + { + std::unique_lock guard(store_client_mutex_); + RAY_ARROW_RETURN_NOT_OK(store_client_.Seal(plasma_id)); + RAY_ARROW_RETURN_NOT_OK(store_client_.Release(plasma_id)); + } + return Status::OK(); } -Status CoreWorkerPlasmaStoreProvider::Get( - const std::vector &ids, int64_t timeout_ms, const TaskID &task_id, - std::vector> *results) { - int64_t batch_size = RayConfig::instance().worker_fetch_request_size(); - (*results).resize(ids.size(), nullptr); - std::unordered_map remaining; +Status CoreWorkerPlasmaStoreProvider::FetchAndGetFromPlasmaStore( + std::unordered_set &remaining, const std::vector &batch_ids, + int64_t timeout_ms, bool fetch_only, const TaskID &task_id, + std::unordered_map> *results, + bool *got_exception) { + RAY_CHECK_OK(raylet_client_->FetchOrReconstruct(batch_ids, fetch_only, task_id)); - // First, attempt to fetch all of the required objects without reconstructing. - for (int64_t start = 0; start < int64_t(ids.size()); start += batch_size) { - int64_t end = std::min(start + batch_size, int64_t(ids.size())); - const std::vector ids_slice(ids.cbegin() + start, ids.cbegin() + end); - RAY_CHECK_OK( - raylet_client_->FetchOrReconstruct(ids_slice, /*fetch_only=*/true, task_id)); - std::vector> results_slice; - RAY_RETURN_NOT_OK(local_store_provider_.Get(ids_slice, 0, task_id, &results_slice)); + std::vector plasma_batch_ids; + plasma_batch_ids.reserve(batch_ids.size()); + for (size_t i = 0; i < batch_ids.size(); i++) { + plasma_batch_ids.push_back(batch_ids[i].ToPlasmaId()); + } + std::vector plasma_results; + { + std::unique_lock guard(store_client_mutex_); + RAY_ARROW_RETURN_NOT_OK( + store_client_.Get(plasma_batch_ids, timeout_ms, &plasma_results)); + } - // Iterate through the results from the local store, adding them to the remaining - // map if they weren't successfully fetched from the local store (are nullptr). - // Keeps track of the locations of the remaining object IDs in the original list. - for (size_t i = 0; i < ids_slice.size(); i++) { - if (results_slice[i] != nullptr) { - (*results)[start + i] = results_slice[i]; - // Terminate early on exception because it'll be raised by the worker anyways. - if (IsException(*results_slice[i])) { - return Status::OK(); - } - } else { - remaining.insert({ids_slice[i], start + i}); + // Add successfully retrieved objects to the result map and remove them from + // the set of IDs to get. + 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]; + const auto result_object = std::make_shared( + std::make_shared(plasma_results[i].data), + std::make_shared(plasma_results[i].metadata)); + (*results)[object_id] = result_object; + remaining.erase(object_id); + if (IsException(*result_object)) { + *got_exception = true; } } } + return Status::OK(); +} + +Status CoreWorkerPlasmaStoreProvider::Get( + const std::unordered_set &object_ids, int64_t timeout_ms, + const TaskID &task_id, + std::unordered_map> *results) { + int64_t batch_size = RayConfig::instance().worker_fetch_request_size(); + bool got_exception = false; + std::vector batch_ids; + std::unordered_set remaining(object_ids.begin(), object_ids.end()); + + // First, attempt to fetch all of the required objects once without reconstructing. + std::vector id_vector(object_ids.begin(), object_ids.end()); + int64_t total_size = static_cast(object_ids.size()); + for (int64_t start = 0; start < total_size; start += batch_size) { + batch_ids.clear(); + for (int64_t i = start; i < batch_size && i < total_size; i++) { + batch_ids.push_back(id_vector[start + i]); + } + RAY_RETURN_NOT_OK(FetchAndGetFromPlasmaStore(remaining, batch_ids, /*timeout_ms=*/0, + /*fetch_only=*/true, task_id, results, + &got_exception)); + } + // If all objects were fetched already, return. - if (remaining.empty()) { + if (remaining.empty() || got_exception) { return Status::OK(); } @@ -60,17 +117,14 @@ Status CoreWorkerPlasmaStoreProvider::Get( bool should_break = false; int64_t remaining_timeout = timeout_ms; while (!remaining.empty() && !should_break) { - std::vector batch_ids; - for (const auto &entry : remaining) { + batch_ids.clear(); + for (const auto &id : remaining) { if (int64_t(batch_ids.size()) == batch_size) { break; } - batch_ids.push_back(entry.first); + batch_ids.push_back(id); } - RAY_CHECK_OK( - raylet_client_->FetchOrReconstruct(batch_ids, /*fetch_only=*/false, task_id)); - int64_t batch_timeout = std::max(RayConfig::instance().get_timeout_milliseconds(), int64_t(10 * batch_ids.size())); if (remaining_timeout >= 0) { @@ -79,26 +133,13 @@ Status CoreWorkerPlasmaStoreProvider::Get( should_break = remaining_timeout <= 0; } - std::vector> batch_results; - RAY_RETURN_NOT_OK( - local_store_provider_.Get(batch_ids, batch_timeout, task_id, &batch_results)); + size_t previous_size = remaining.size(); + RAY_RETURN_NOT_OK(FetchAndGetFromPlasmaStore(remaining, batch_ids, batch_timeout, + /*fetch_only=*/false, task_id, results, + &got_exception)); + should_break = should_break || got_exception; - // Add successfully retrieved objects to the result list and remove them from - // remaining. - uint64_t successes = 0; - for (size_t i = 0; i < batch_results.size(); i++) { - if (batch_results[i] != nullptr) { - successes++; - const auto &object_id = batch_ids[i]; - (*results)[remaining[object_id]] = batch_results[i]; - remaining.erase(object_id); - if (IsException(*batch_results[i])) { - should_break = true; - } - } - } - - if (successes < batch_ids.size()) { + if ((previous_size - remaining.size()) < batch_ids.size()) { unsuccessful_attempts++; WarnIfAttemptedTooManyTimes(unsuccessful_attempts, remaining); } @@ -153,12 +194,12 @@ bool CoreWorkerPlasmaStoreProvider::IsException(const RayObject &object) { } void CoreWorkerPlasmaStoreProvider::WarnIfAttemptedTooManyTimes( - int num_attempts, const std::unordered_map &remaining) { + int num_attempts, const std::unordered_set &remaining) { if (num_attempts % RayConfig::instance().object_store_get_warn_per_num_attempts() == 0) { std::ostringstream oss; size_t printed = 0; - for (auto &entry : remaining) { + for (auto &id : remaining) { if (printed >= RayConfig::instance().object_store_get_max_ids_to_print_in_warning()) { break; @@ -166,7 +207,7 @@ void CoreWorkerPlasmaStoreProvider::WarnIfAttemptedTooManyTimes( if (printed > 0) { oss << ", "; } - oss << entry.first.Hex(); + oss << id.Hex(); } if (printed < remaining.size()) { oss << ", etc"; 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 4fa5537bc..04d666084 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.h +++ b/src/ray/core_worker/store_provider/plasma_store_provider.h @@ -6,7 +6,6 @@ #include "ray/common/id.h" #include "ray/common/status.h" #include "ray/core_worker/common.h" -#include "ray/core_worker/store_provider/local_plasma_provider.h" #include "ray/core_worker/store_provider/store_provider.h" #include "ray/raylet/raylet_client.h" @@ -21,47 +20,46 @@ class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider { CoreWorkerPlasmaStoreProvider(const std::string &store_socket, std::unique_ptr &raylet_client); - /// Put an object with specified ID into object store. - /// - /// \param[in] object The ray object. - /// \param[in] object_id Object ID specified by user. - /// \return Status. + /// See `CoreWorkerStoreProvider::Put` for semantics. Status Put(const RayObject &object, const ObjectID &object_id) override; - /// Get a list of objects from the object store. - /// - /// \param[in] ids IDs of the objects to get. - /// \param[in] timeout_ms Timeout in milliseconds, wait infinitely if it's negative. - /// \param[in] task_id ID for the current task. - /// \param[out] results Result list of objects data. - /// \return Status. - Status Get(const std::vector &ids, int64_t timeout_ms, const TaskID &task_id, - std::vector> *results) override; + /// See `CoreWorkerStoreProvider::Get` for semantics. + Status Get(const std::unordered_set &object_ids, int64_t timeout_ms, + const TaskID &task_id, + std::unordered_map> *results) override; - /// Wait for a list of objects to appear in the object store. - /// - /// \param[in] IDs of the objects to wait for. - /// \param[in] num_returns Number of objects that should appear. - /// \param[in] timeout_ms Timeout in milliseconds, wait infinitely if it's negative. - /// \param[in] task_id ID for the current task. - /// \param[out] results A bitset that indicates each object has appeared or not. - /// \return Status. + /// See `CoreWorkerStoreProvider::Wait` for semantics. Status Wait(const std::vector &object_ids, int num_objects, int64_t timeout_ms, const TaskID &task_id, std::vector *results) override; - /// Delete a list of objects from the object store. - /// - /// \param[in] object_ids IDs of the objects to delete. - /// \param[in] local_only Whether only delete the objects in local node, or all nodes in - /// the cluster. - /// \param[in] delete_creating_tasks Whether also delete the tasks that - /// created these objects. - /// \return Status. + /// See `CoreWorkerStoreProvider::Delete` for semantics. Status Delete(const std::vector &object_ids, bool local_only = true, bool delete_creating_tasks = false) override; private: + /// Ask the raylet to fetch a set of objects and then attempt to get them + /// from the local plasma store. Successfully fetched objects will be removed + /// from the input set of remaining IDs and added to the results map. + /// + /// \param[in/out] remaining IDs of the remaining objects to get. + /// \param[in] batch_ids IDs of the objects to get. + /// \param[in] timeout_ms Timeout in milliseconds. + /// \param[in] fetch_only Whether the raylet should only fetch or also attempt to + /// reconstruct objects. + /// \param[in] task_id The current TaskID. + /// \param[out] results Map of objects to write results into. This method will only + /// add to this map, not clear or remove from it, so the caller can pass in a non-empty + /// map. + /// \param[out] got_exception Whether any of the fetched objects contained an + /// exception. + /// \return Status. + Status FetchAndGetFromPlasmaStore( + std::unordered_set &remaining, const std::vector &batch_ids, + int64_t timeout_ms, bool fetch_only, const TaskID &task_id, + std::unordered_map> *results, + bool *got_exception); + /// Whether the buffer represents an exception object. /// /// \param[in] object Object data. @@ -73,14 +71,12 @@ class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider { /// /// \param[in] num_attemps The number of attempted times. /// \param[in] remaining The remaining objects. - static void WarnIfAttemptedTooManyTimes( - int num_attempts, const std::unordered_map &remaining); + static void WarnIfAttemptedTooManyTimes(int num_attempts, + const std::unordered_set &remaining); - /// local plasma store provider. - CoreWorkerLocalPlasmaStoreProvider local_store_provider_; - - /// Raylet client. std::unique_ptr &raylet_client_; + plasma::PlasmaClient store_client_; + std::mutex store_client_mutex_; }; } // namespace ray diff --git a/src/ray/core_worker/store_provider/store_provider.h b/src/ray/core_worker/store_provider/store_provider.h index 63dbbca64..1c3474fc4 100644 --- a/src/ray/core_worker/store_provider/store_provider.h +++ b/src/ray/core_worker/store_provider/store_provider.h @@ -8,7 +8,7 @@ namespace ray { -/// Binary representation of ray object. +/// Binary representation of a ray object. class RayObject { public: /// Create a ray object instance. @@ -75,20 +75,22 @@ class CoreWorkerStoreProvider { /// \return Status. virtual Status Put(const RayObject &object, const ObjectID &object_id) = 0; - /// Get a list of objects from the object store. + /// Get a set of objects from the object store. /// - /// \param[in] ids IDs of the objects to get. + /// \param[in] object_ids IDs of the objects to get. /// \param[in] timeout_ms Timeout in milliseconds, wait infinitely if it's negative. /// \param[in] task_id ID for the current task. - /// \param[out] results Result list of objects data. + /// \param[out] results Map of objects to write results into. Get will only add to this + /// map, not clear or remove from it, so the caller can pass in a non-empty map. /// \return Status. - virtual Status Get(const std::vector &ids, int64_t timeout_ms, - const TaskID &task_id, - std::vector> *results) = 0; + virtual Status Get( + const std::unordered_set &object_ids, int64_t timeout_ms, + const TaskID &task_id, + std::unordered_map> *results) = 0; /// Wait for a list of objects to appear in the object store. /// - /// \param[in] IDs of the objects to wait for. + /// \param[in] object_ids IDs of the objects to wait for. /// \param[in] num_objects Number of objects that should appear before returning. /// \param[in] timeout_ms Timeout in milliseconds, wait infinitely if it's negative. /// \param[in] task_id ID for the current task. diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc index 99a7f5ca8..bc48162f1 100644 --- a/src/ray/core_worker/test/core_worker_test.cc +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -7,7 +7,6 @@ #include "ray/core_worker/core_worker.h" #include "ray/core_worker/transport/direct_actor_transport.h" -#include "ray/core_worker/store_provider/local_plasma_provider.h" #include "ray/core_worker/store_provider/memory_store_provider.h" #include "ray/raylet/raylet_client.h" @@ -468,10 +467,6 @@ void CoreWorkerTest::TestStoreProvider(StoreProviderType type) { std::shared_ptr memory_store; switch (type) { - case StoreProviderType::LOCAL_PLASMA: - provider_ptr = std::unique_ptr( - new CoreWorkerLocalPlasmaStoreProvider(raylet_store_socket_names_[0])); - break; case StoreProviderType::MEMORY: memory_store = std::make_shared(); provider_ptr = std::unique_ptr( @@ -522,19 +517,20 @@ void CoreWorkerTest::TestStoreProvider(StoreProviderType type) { ASSERT_EQ(wait_results, std::vector({true, true, true, true, false})); // Test Get(). - std::vector> results; - RAY_CHECK_OK(provider.Get(ids_with_duplicate, -1, RandomTaskId(), &results)); + std::unordered_map> results; + std::unordered_set ids_set(ids.begin(), ids.end()); + RAY_CHECK_OK(provider.Get(ids_set, -1, RandomTaskId(), &results)); - ASSERT_EQ(results.size(), ids_with_duplicate.size()); - for (size_t i = 0; i < ids_with_duplicate.size(); i++) { - const auto &expected = buffers[i % ids.size()]; - ASSERT_EQ(results[i]->GetData()->Size(), expected.GetData()->Size()); - ASSERT_EQ(memcmp(results[i]->GetData()->Data(), expected.GetData()->Data(), + ASSERT_EQ(results.size(), ids.size()); + for (size_t i = 0; i < ids.size(); i++) { + const auto &expected = buffers[i]; + ASSERT_EQ(results[ids[i]]->GetData()->Size(), expected.GetData()->Size()); + ASSERT_EQ(memcmp(results[ids[i]]->GetData()->Data(), expected.GetData()->Data(), expected.GetData()->Size()), 0); - ASSERT_EQ(results[i]->GetMetadata()->Size(), expected.GetMetadata()->Size()); - ASSERT_EQ(memcmp(results[i]->GetMetadata()->Data(), expected.GetMetadata()->Data(), - expected.GetMetadata()->Size()), + ASSERT_EQ(results[ids[i]]->GetMetadata()->Size(), expected.GetMetadata()->Size()); + ASSERT_EQ(memcmp(results[ids[i]]->GetMetadata()->Data(), + expected.GetMetadata()->Data(), expected.GetMetadata()->Size()), 0); } @@ -545,10 +541,8 @@ void CoreWorkerTest::TestStoreProvider(StoreProviderType type) { RAY_CHECK_OK(provider.Delete(ids, true, false)); usleep(200 * 1000); - RAY_CHECK_OK(provider.Get(ids, 0, RandomTaskId(), &results)); - ASSERT_EQ(results.size(), 2); - ASSERT_TRUE(!results[0]); - ASSERT_TRUE(!results[1]); + RAY_CHECK_OK(provider.Get(ids_set, 0, RandomTaskId(), &results)); + ASSERT_EQ(results.size(), 0); // Test Wait() with objects which will become ready later. std::vector unready_ids(buffers.size()); @@ -789,8 +783,7 @@ TEST_F(SingleNodeTest, TestObjectInterface) { // Test Get(). std::vector> results; RAY_CHECK_OK(core_worker.Objects().Get(ids, -1, &results)); - - ASSERT_EQ(results.size(), 2); + ASSERT_EQ(results.size(), ids.size()); for (size_t i = 0; i < ids.size(); i++) { ASSERT_EQ(*results[i]->GetData(), *buffers[i].GetData()); ASSERT_EQ(*results[i]->GetMetadata(), *buffers[i].GetMetadata());