mirror of
https://github.com/wassname/ray.git
synced 2026-07-10 20:36:41 +08:00
Clean up Wait() and Get() in the core worker (#5556)
This commit is contained in:
committed by
Philipp Moritz
parent
4ed6ee0b1e
commit
0c68b4cc30
@@ -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 };
|
||||
|
||||
|
||||
@@ -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_));
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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());
|
||||
|
||||
Reference in New Issue
Block a user