Clean up Wait() and Get() in the core worker (#5556)

This commit is contained in:
Edward Oakes
2019-09-03 14:45:15 -07:00
committed by Philipp Moritz
parent 4ed6ee0b1e
commit 0c68b4cc30
11 changed files with 189 additions and 333 deletions
+1 -1
View File
@@ -64,7 +64,7 @@ class TaskArg {
const std::shared_ptr<Buffer> data_;
};
enum class StoreProviderType { LOCAL_PLASMA, PLASMA, MEMORY };
enum class StoreProviderType { PLASMA, MEMORY };
enum class TaskTransportType { RAYLET, DIRECT_ACTOR };
+10 -30
View File
@@ -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<CoreWorkerMemoryStore>()) {
AddStoreProvider(StoreProviderType::LOCAL_PLASMA);
AddStoreProvider(StoreProviderType::PLASMA);
AddStoreProvider(StoreProviderType::MEMORY);
}
@@ -76,17 +74,17 @@ Status CoreWorkerObjectInterface::Get(const std::vector<ObjectID> &ids,
object_ids_per_store_provider;
GroupObjectIdsByStoreProvider(ids, &object_ids_per_store_provider);
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> objects;
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> 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<StoreProviderType, std::reference_wrapper<std::unordered_set<ObjectID>>>>
std::list<std::pair<StoreProviderType,
std::reference_wrapper<const std::unordered_set<ObjectID>>>>
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<ObjectID> &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<ObjectID> &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<ObjectID> &object_ids,
int64_t timeout_ms,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *results) {
std::vector<ObjectID> ids(object_ids.begin(), object_ids.end());
if (!ids.empty()) {
std::vector<std::shared_ptr<RayObject>> 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<CoreWorkerStoreProvider> CoreWorkerObjectInterface::CreateStoreProvider(
StoreProviderType type) const {
switch (type) {
case StoreProviderType::LOCAL_PLASMA:
return std::unique_ptr<CoreWorkerStoreProvider>(
new CoreWorkerLocalPlasmaStoreProvider(store_socket_));
case StoreProviderType::PLASMA:
return std::unique_ptr<CoreWorkerStoreProvider>(
new CoreWorkerPlasmaStoreProvider(store_socket_, raylet_client_));
-12
View File
@@ -86,18 +86,6 @@ class CoreWorkerObjectInterface {
&ids_per_provider,
int64_t timeout_ms, int *num_objects, std::vector<bool> *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<ObjectID> &object_ids,
int64_t timeout_ms,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *results);
/// Helper function to wait a list of objects from a specific store provider.
///
/// \param[in] type The type of store provider to use.
@@ -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<arrow::Buffer> out_buffer;
{
std::unique_lock<std::mutex> 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<std::mutex> 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<ObjectID> &object_ids, int64_t timeout_ms, const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *results) {
std::vector<plasma::ObjectID> plasma_ids;
plasma_ids.reserve(object_ids.size());
for (const auto &object_id : object_ids) {
plasma_ids.push_back(object_id.ToPlasmaId());
}
std::vector<plasma::ObjectBuffer> object_buffers;
{
std::unique_lock<std::mutex> 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<RayObject>(
std::make_shared<PlasmaBuffer>(object_buffers[i].data),
std::make_shared<PlasmaBuffer>(object_buffers[i].metadata));
}
}
return Status::OK();
}
Status CoreWorkerLocalPlasmaStoreProvider::Wait(const std::vector<ObjectID> &object_ids,
int num_objects, int64_t timeout_ms,
const TaskID &task_id,
std::vector<bool> *results) {
if (num_objects != static_cast<int>(object_ids.size())) {
return Status::Invalid("num_objects should equal to number of items in object_ids");
}
std::vector<std::shared_ptr<RayObject>> 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<ObjectID> &object_ids,
bool local_only,
bool delete_creating_tasks) {
std::vector<plasma::ObjectID> 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<std::mutex> guard(store_client_mutex_);
RAY_ARROW_RETURN_NOT_OK(store_client_.Delete(plasma_ids));
return Status::OK();
}
} // namespace ray
@@ -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<ObjectID> &ids, int64_t timeout_ms, const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *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<ObjectID> &object_ids, int num_objects,
int64_t timeout_ms, const TaskID &task_id,
std::vector<bool> *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<ObjectID> &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
@@ -22,9 +22,20 @@ Status CoreWorkerMemoryStoreProvider::Put(const RayObject &object,
}
Status CoreWorkerMemoryStoreProvider::Get(
const std::vector<ObjectID> &object_ids, int64_t timeout_ms, const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *results) {
return store_->Get(object_ids, object_ids.size(), timeout_ms, true, results);
const std::unordered_set<ObjectID> &object_ids, int64_t timeout_ms,
const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *results) {
const std::vector<ObjectID> id_vector(object_ids.begin(), object_ids.end());
std::vector<std::shared_ptr<RayObject>> 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<ObjectID> &object_ids,
@@ -34,15 +45,14 @@ Status CoreWorkerMemoryStoreProvider::Wait(const std::vector<ObjectID> &object_i
(*results).resize(object_ids.size(), false);
std::vector<std::shared_ptr<RayObject>> 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<ObjectID> &object_ids,
@@ -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<ObjectID> &ids, int64_t timeout_ms, const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *results) override;
Status Get(const std::unordered_set<ObjectID> &object_ids, int64_t timeout_ms,
const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *results) override;
/// See `CoreWorkerStoreProvider::Wait` for semantics.
/// Note that `num_objects` must equal to number of items in `object_ids`.
@@ -9,47 +9,104 @@ namespace ray {
CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider(
const std::string &store_socket, std::unique_ptr<RayletClient> &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<arrow::Buffer> out_buffer;
{
std::unique_lock<std::mutex> 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<std::mutex> 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<ObjectID> &ids, int64_t timeout_ms, const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *results) {
int64_t batch_size = RayConfig::instance().worker_fetch_request_size();
(*results).resize(ids.size(), nullptr);
std::unordered_map<ObjectID, int> remaining;
Status CoreWorkerPlasmaStoreProvider::FetchAndGetFromPlasmaStore(
std::unordered_set<ObjectID> &remaining, const std::vector<ObjectID> &batch_ids,
int64_t timeout_ms, bool fetch_only, const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *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<ObjectID> ids_slice(ids.cbegin() + start, ids.cbegin() + end);
RAY_CHECK_OK(
raylet_client_->FetchOrReconstruct(ids_slice, /*fetch_only=*/true, task_id));
std::vector<std::shared_ptr<RayObject>> results_slice;
RAY_RETURN_NOT_OK(local_store_provider_.Get(ids_slice, 0, task_id, &results_slice));
std::vector<plasma::ObjectID> 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::ObjectBuffer> plasma_results;
{
std::unique_lock<std::mutex> 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<RayObject>(
std::make_shared<PlasmaBuffer>(plasma_results[i].data),
std::make_shared<PlasmaBuffer>(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<ObjectID> &object_ids, int64_t timeout_ms,
const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *results) {
int64_t batch_size = RayConfig::instance().worker_fetch_request_size();
bool got_exception = false;
std::vector<ObjectID> batch_ids;
std::unordered_set<ObjectID> remaining(object_ids.begin(), object_ids.end());
// First, attempt to fetch all of the required objects once without reconstructing.
std::vector<ObjectID> id_vector(object_ids.begin(), object_ids.end());
int64_t total_size = static_cast<int64_t>(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<ObjectID> 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<std::shared_ptr<RayObject>> 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<ObjectID, int> &remaining) {
int num_attempts, const std::unordered_set<ObjectID> &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";
@@ -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<RayletClient> &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<ObjectID> &ids, int64_t timeout_ms, const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *results) override;
/// See `CoreWorkerStoreProvider::Get` for semantics.
Status Get(const std::unordered_set<ObjectID> &object_ids, int64_t timeout_ms,
const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *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<ObjectID> &object_ids, int num_objects,
int64_t timeout_ms, const TaskID &task_id,
std::vector<bool> *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<ObjectID> &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<ObjectID> &remaining, const std::vector<ObjectID> &batch_ids,
int64_t timeout_ms, bool fetch_only, const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *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<ObjectID, int> &remaining);
static void WarnIfAttemptedTooManyTimes(int num_attempts,
const std::unordered_set<ObjectID> &remaining);
/// local plasma store provider.
CoreWorkerLocalPlasmaStoreProvider local_store_provider_;
/// Raylet client.
std::unique_ptr<RayletClient> &raylet_client_;
plasma::PlasmaClient store_client_;
std::mutex store_client_mutex_;
};
} // namespace ray
@@ -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<ObjectID> &ids, int64_t timeout_ms,
const TaskID &task_id,
std::vector<std::shared_ptr<RayObject>> *results) = 0;
virtual Status Get(
const std::unordered_set<ObjectID> &object_ids, int64_t timeout_ms,
const TaskID &task_id,
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> *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.
+14 -21
View File
@@ -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<CoreWorkerMemoryStore> memory_store;
switch (type) {
case StoreProviderType::LOCAL_PLASMA:
provider_ptr = std::unique_ptr<CoreWorkerStoreProvider>(
new CoreWorkerLocalPlasmaStoreProvider(raylet_store_socket_names_[0]));
break;
case StoreProviderType::MEMORY:
memory_store = std::make_shared<CoreWorkerMemoryStore>();
provider_ptr = std::unique_ptr<CoreWorkerStoreProvider>(
@@ -522,19 +517,20 @@ void CoreWorkerTest::TestStoreProvider(StoreProviderType type) {
ASSERT_EQ(wait_results, std::vector<bool>({true, true, true, true, false}));
// Test Get().
std::vector<std::shared_ptr<RayObject>> results;
RAY_CHECK_OK(provider.Get(ids_with_duplicate, -1, RandomTaskId(), &results));
std::unordered_map<ObjectID, std::shared_ptr<RayObject>> results;
std::unordered_set<ObjectID> 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<ObjectID> unready_ids(buffers.size());
@@ -789,8 +783,7 @@ TEST_F(SingleNodeTest, TestObjectInterface) {
// Test Get().
std::vector<std::shared_ptr<RayObject>> 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());