diff --git a/src/ray/common/buffer.h b/src/ray/common/buffer.h index 4340c74a8..d14cb3fbb 100644 --- a/src/ray/common/buffer.h +++ b/src/ray/common/buffer.h @@ -30,7 +30,14 @@ class Buffer { /// Represents a byte buffer in local memory. class LocalMemoryBuffer : public Buffer { public: - LocalMemoryBuffer(uint8_t *data, size_t size) : data_(data), size_(size) {} + LocalMemoryBuffer(uint8_t *data, size_t size, bool should_copy = false) + : data_(data), size_(size) { + if (should_copy) { + buffer_.insert(buffer_.end(), data, data + size); + data_ = buffer_.data(); + size_ = buffer_.size(); + } + } uint8_t *Data() const override { return data_; } @@ -43,6 +50,8 @@ class LocalMemoryBuffer : public Buffer { uint8_t *data_; /// Size of the buffer. size_t size_; + /// This is only valid when `should_copy` is true. + std::vector buffer_; }; /// Represents a byte buffer for plasma object. diff --git a/src/ray/core_worker/common.h b/src/ray/core_worker/common.h index d265aa536..ef3188266 100644 --- a/src/ray/core_worker/common.h +++ b/src/ray/core_worker/common.h @@ -7,6 +7,7 @@ #include "ray/common/id.h" #include "ray/common/task/task_spec.h" #include "ray/raylet/raylet_client.h" +#include "ray/util/util.h" namespace ray { using WorkerType = rpc::WorkerType; @@ -73,7 +74,7 @@ struct TaskInfo { const TaskType task_type; }; -enum class StoreProviderType { PLASMA }; +enum class StoreProviderType { LOCAL_PLASMA, PLASMA }; enum class TaskTransportType { RAYLET }; diff --git a/src/ray/core_worker/mock_worker.cc b/src/ray/core_worker/mock_worker.cc index d77931c87..7117f4ec4 100644 --- a/src/ray/core_worker/mock_worker.cc +++ b/src/ray/core_worker/mock_worker.cc @@ -21,7 +21,7 @@ class MockWorker { MockWorker(const std::string &store_socket, const std::string &raylet_socket) : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket, JobID::JobID::FromInt(1), - std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4)) {} + std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4, _5)) {} void Run() { // Start executing tasks. @@ -31,7 +31,8 @@ class MockWorker { private: Status ExecuteTask(const RayFunction &ray_function, const std::vector> &args, - const TaskInfo &task_info, int num_returns) { + const TaskInfo &task_info, int num_returns, + std::vector> *results) { // Note that this doesn't include dummy object id. RAY_CHECK(num_returns >= 0); @@ -41,14 +42,12 @@ class MockWorker { auto &data = arg->GetData(); buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size()); } - - auto return_value = RayObject( - std::make_shared(buffer.data(), buffer.size()), nullptr); + auto memory_buffer = + std::make_shared(buffer.data(), buffer.size(), true); // Write the merged content to each of return ids. for (int i = 0; i < num_returns; i++) { - ObjectID id = ObjectID::ForTaskReturn(task_info.task_id, i + 1); - RAY_CHECK_OK(worker_.Objects().Put(return_value, id)); + results->push_back(std::make_shared(memory_buffer, nullptr)); } return Status::OK(); } diff --git a/src/ray/core_worker/object_interface.cc b/src/ray/core_worker/object_interface.cc index 8a4eb4642..ea29d7d62 100644 --- a/src/ray/core_worker/object_interface.cc +++ b/src/ray/core_worker/object_interface.cc @@ -1,5 +1,6 @@ #include "ray/core_worker/object_interface.h" #include "ray/common/ray_config.h" +#include "ray/core_worker/store_provider/local_plasma_provider.h" #include "ray/core_worker/store_provider/plasma_store_provider.h" namespace ray { @@ -7,11 +8,11 @@ namespace ray { CoreWorkerObjectInterface::CoreWorkerObjectInterface( WorkerContext &worker_context, std::unique_ptr &raylet_client, const std::string &store_socket) - : worker_context_(worker_context), raylet_client_(raylet_client) { - store_providers_.emplace( - static_cast(StoreProviderType::PLASMA), - std::unique_ptr( - new CoreWorkerPlasmaStoreProvider(store_socket, raylet_client_))); + : worker_context_(worker_context), + raylet_client_(raylet_client), + store_socket_(store_socket) { + AddStoreProvider(StoreProviderType::LOCAL_PLASMA); + AddStoreProvider(StoreProviderType::PLASMA); } Status CoreWorkerObjectInterface::Put(const RayObject &object, ObjectID *object_id) { @@ -23,30 +24,48 @@ Status CoreWorkerObjectInterface::Put(const RayObject &object, ObjectID *object_ Status CoreWorkerObjectInterface::Put(const RayObject &object, const ObjectID &object_id) { - auto type = static_cast(StoreProviderType::PLASMA); - return store_providers_[type]->Put(object, object_id); + return store_providers_[StoreProviderType::PLASMA]->Put(object, object_id); } Status CoreWorkerObjectInterface::Get(const std::vector &ids, int64_t timeout_ms, std::vector> *results) { - auto type = static_cast(StoreProviderType::PLASMA); - return store_providers_[type]->Get(ids, timeout_ms, worker_context_.GetCurrentTaskID(), - results); + return store_providers_[StoreProviderType::PLASMA]->Get( + ids, timeout_ms, worker_context_.GetCurrentTaskID(), results); } Status CoreWorkerObjectInterface::Wait(const std::vector &object_ids, int num_objects, int64_t timeout_ms, std::vector *results) { - auto type = static_cast(StoreProviderType::PLASMA); - return store_providers_[type]->Wait(object_ids, num_objects, timeout_ms, - worker_context_.GetCurrentTaskID(), results); + return store_providers_[StoreProviderType::PLASMA]->Wait( + object_ids, num_objects, timeout_ms, worker_context_.GetCurrentTaskID(), results); } Status CoreWorkerObjectInterface::Delete(const std::vector &object_ids, bool local_only, bool delete_creating_tasks) { - auto type = static_cast(StoreProviderType::PLASMA); - return store_providers_[type]->Delete(object_ids, local_only, delete_creating_tasks); + return store_providers_[StoreProviderType::PLASMA]->Delete(object_ids, local_only, + delete_creating_tasks); +} + +void CoreWorkerObjectInterface::AddStoreProvider(StoreProviderType type) { + store_providers_.emplace(type, CreateStoreProvider(type)); +} + +std::unique_ptr CoreWorkerObjectInterface::CreateStoreProvider( + StoreProviderType type) const { + switch (type) { + case StoreProviderType::LOCAL_PLASMA: + return std::unique_ptr( + new CoreWorkerLocalPlasmaStoreProvider(store_socket_)); + break; + case StoreProviderType::PLASMA: + return std::unique_ptr( + new CoreWorkerPlasmaStoreProvider(store_socket_, raylet_client_)); + break; + default: + RAY_LOG(FATAL) << "unknown store provider type " << static_cast(type); + break; + } } } // namespace ray diff --git a/src/ray/core_worker/object_interface.h b/src/ray/core_worker/object_interface.h index f1d66e803..cd9e461ee 100644 --- a/src/ray/core_worker/object_interface.h +++ b/src/ray/core_worker/object_interface.h @@ -66,13 +66,26 @@ class CoreWorkerObjectInterface { bool delete_creating_tasks); private: + /// Create a new store provider for the specified type on demand. + std::unique_ptr CreateStoreProvider( + StoreProviderType type) const; + + /// Add a store provider for the specified type. + void AddStoreProvider(StoreProviderType type); + /// Reference to the parent CoreWorker's context. WorkerContext &worker_context_; /// Reference to the parent CoreWorker's raylet client. std::unique_ptr &raylet_client_; + /// Store socket name. + std::string store_socket_; + /// All the store providers supported. - std::unordered_map> store_providers_; + EnumUnorderedMap> + store_providers_; + + friend class CoreWorkerTaskInterface; }; } // namespace ray diff --git a/src/ray/core_worker/store_provider/local_plasma_provider.cc b/src/ray/core_worker/store_provider/local_plasma_provider.cc new file mode 100644 index 000000000..758ebd72a --- /dev/null +++ b/src/ray/core_worker/store_provider/local_plasma_provider.cc @@ -0,0 +1,104 @@ +#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) { + auto status = store_client_.Connect(store_socket); + if (!status.ok()) { + RAY_LOG(ERROR) << "Connecting plasma store failed when trying to construct" + << " core worker: " << status.message(); + throw std::runtime_error(status.message()); + } +} + +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_); + RAY_ARROW_RETURN_NOT_OK(store_client_.Create( + plasma_id, data ? data->Size() : 0, metadata ? metadata->Data() : nullptr, + metadata ? metadata->Size() : 0, &out_buffer)); + } + + 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 != 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]->GetData() != 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 new file mode 100644 index 000000000..d67b916b9 --- /dev/null +++ b/src/ray/core_worker/store_provider/local_plasma_provider.h @@ -0,0 +1,49 @@ +#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/plasma_store_provider.cc b/src/ray/core_worker/store_provider/plasma_store_provider.cc index 5a59420a5..dc6d415ee 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.cc +++ b/src/ray/core_worker/store_provider/plasma_store_provider.cc @@ -9,36 +9,11 @@ namespace ray { CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider( const std::string &store_socket, std::unique_ptr &raylet_client) - : raylet_client_(raylet_client) { - auto status = store_client_.Connect(store_socket); - if (!status.ok()) { - RAY_LOG(ERROR) << "Connecting plasma store failed when trying to construct" - << " core worker: " << status.message(); - throw std::runtime_error(status.message()); - } -} + : local_store_provider_(store_socket), raylet_client_(raylet_client) {} Status CoreWorkerPlasmaStoreProvider::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_); - RAY_ARROW_RETURN_NOT_OK(store_client_.Create( - plasma_id, data->Size(), metadata ? metadata->Data() : nullptr, - metadata ? metadata->Size() : 0, &out_buffer)); - } - - 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(); + return local_store_provider_.Put(object, object_id); } Status CoreWorkerPlasmaStoreProvider::Get( @@ -84,25 +59,16 @@ Status CoreWorkerPlasmaStoreProvider::Get( get_timeout = RayConfig::instance().get_timeout_milliseconds(); } - std::vector plasma_ids; - for (const auto &id : unready_ids) { - plasma_ids.push_back(id.ToPlasmaId()); - } + std::vector> result_objects; + RAY_RETURN_NOT_OK( + local_store_provider_.Get(unready_ids, get_timeout, task_id, &result_objects)); - std::vector object_buffers; - { - std::unique_lock guard(store_client_mutex_); - auto status = store_client_.Get(plasma_ids, get_timeout, &object_buffers); - } - - for (size_t i = 0; i < object_buffers.size(); i++) { - if (object_buffers[i].data != nullptr) { + for (size_t i = 0; i < result_objects.size(); i++) { + if (result_objects[i] != nullptr) { const auto &object_id = unready_ids[i]; - (*results)[unready[object_id]] = std::make_shared( - std::make_shared(object_buffers[i].data), - std::make_shared(object_buffers[i].metadata)); + (*results)[unready[object_id]] = result_objects[i]; unready.erase(object_id); - if (IsException(object_buffers[i])) { + if (IsException(*result_objects[i])) { should_break = true; } } @@ -148,9 +114,10 @@ Status CoreWorkerPlasmaStoreProvider::Delete(const std::vector &object return raylet_client_->FreeObjects(object_ids, local_only, delete_creating_tasks); } -bool CoreWorkerPlasmaStoreProvider::IsException(const plasma::ObjectBuffer &buffer) { +bool CoreWorkerPlasmaStoreProvider::IsException(const RayObject &object) { // TODO (kfstorm): metadata should be structured. - const std::string metadata = buffer.metadata->ToString(); + const std::string metadata(reinterpret_cast(object.GetMetadata()->Data()), + object.GetMetadata()->Size()); const auto error_type_descriptor = ray::rpc::ErrorType_descriptor(); for (int i = 0; i < error_type_descriptor->value_count(); i++) { const auto error_type_number = error_type_descriptor->value(i)->number(); 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 89ecb2ea2..1ef6a8b1d 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.h +++ b/src/ray/core_worker/store_provider/plasma_store_provider.h @@ -6,6 +6,7 @@ #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" @@ -55,16 +56,17 @@ class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider { /// \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. - Status Delete(const std::vector &object_ids, bool local_only, - bool delete_creating_tasks) override; + /// created these objects. + /// \return Status. + Status Delete(const std::vector &object_ids, bool local_only = true, + bool delete_creating_tasks = false) override; private: /// Whether the buffer represents an exception object. /// - /// \param[in] buffer the object buffer. + /// \param[in] object Object data. /// \return Whether it represents an exception object. - static bool IsException(const plasma::ObjectBuffer &buffer); + static bool IsException(const RayObject &object); /// Print a warning if we've attempted too many times, but some objects are still /// unavailable. @@ -74,11 +76,8 @@ class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider { static void WarnIfAttemptedTooManyTimes( int num_attempts, const std::unordered_map &unready); - /// Plasma store client. - plasma::PlasmaClient store_client_; - - /// Mutex to protect store_client_. - std::mutex store_client_mutex_; + /// local plasma store provider. + CoreWorkerLocalPlasmaStoreProvider local_store_provider_; /// Raylet client. std::unique_ptr &raylet_client_; diff --git a/src/ray/core_worker/store_provider/store_provider.h b/src/ray/core_worker/store_provider/store_provider.h index fdba0d887..2533fe8b9 100644 --- a/src/ray/core_worker/store_provider/store_provider.h +++ b/src/ray/core_worker/store_provider/store_provider.h @@ -79,8 +79,8 @@ class CoreWorkerStoreProvider { /// \param[in] delete_creating_tasks Whether also delete the tasks that /// created these objects. /// \return Status. - virtual Status Delete(const std::vector &object_ids, bool local_only, - bool delete_creating_tasks) = 0; + virtual Status Delete(const std::vector &object_ids, bool local_only = true, + bool delete_creating_tasks = false) = 0; }; } // namespace ray diff --git a/src/ray/core_worker/task_execution.cc b/src/ray/core_worker/task_execution.cc index 5f939d3cf..c3ddee166 100644 --- a/src/ray/core_worker/task_execution.cc +++ b/src/ray/core_worker/task_execution.cc @@ -16,17 +16,19 @@ CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface( RAY_CHECK(execution_callback_ != nullptr); auto func = std::bind(&CoreWorkerTaskExecutionInterface::ExecuteTask, this, - std::placeholders::_1); + std::placeholders::_1, std::placeholders::_2); task_receivers_.emplace( - static_cast(TaskTransportType::RAYLET), + TaskTransportType::RAYLET, std::unique_ptr(new CoreWorkerRayletTaskReceiver( - raylet_client, main_service_, worker_server_, func))); + raylet_client, object_interface_, main_service_, worker_server_, func))); // Start RPC server after all the task receivers are properly initialized. worker_server_.Run(); } -Status CoreWorkerTaskExecutionInterface::ExecuteTask(const TaskSpecification &task_spec) { +Status CoreWorkerTaskExecutionInterface::ExecuteTask( + const TaskSpecification &task_spec, + std::vector> *results) { worker_context_.SetCurrentTask(task_spec); RayFunction func{task_spec.GetLanguage(), task_spec.FunctionDescriptor()}; @@ -48,11 +50,12 @@ Status CoreWorkerTaskExecutionInterface::ExecuteTask(const TaskSpecification &ta auto num_returns = task_spec.NumReturns(); if (task_spec.IsActorCreationTask() || task_spec.IsActorTask()) { RAY_CHECK(num_returns > 0); - // Decrease to account for the dummy object id. + // Decrease to account for the dummy object id, this logic only + // applies to task submitted via raylet. num_returns--; } - auto status = execution_callback_(func, args, task_info, num_returns); + auto status = execution_callback_(func, args, task_info, num_returns, results); // TODO(zhijunfu): // 1. Check and handle failure. // 2. Save or load checkpoint. diff --git a/src/ray/core_worker/task_execution.h b/src/ray/core_worker/task_execution.h index 599bd71d7..c8b164642 100644 --- a/src/ray/core_worker/task_execution.h +++ b/src/ray/core_worker/task_execution.h @@ -27,11 +27,13 @@ class CoreWorkerTaskExecutionInterface { /// /// \param ray_function[in] Information about the function to execute. /// \param args[in] Arguments of the task. + /// \param task_info[in] Information of the task to execute. + /// \param results[out] Results of the task execution. /// \return Status. - using TaskExecutor = - std::function> &args, - const TaskInfo &task_info, int num_returns)>; + using TaskExecutor = std::function> &args, const TaskInfo &task_info, + int num_returns, std::vector> *results)>; CoreWorkerTaskExecutionInterface(WorkerContext &worker_context, std::unique_ptr &raylet_client, @@ -55,7 +57,12 @@ class CoreWorkerTaskExecutionInterface { std::vector> *args); /// Execute a task. - Status ExecuteTask(const TaskSpecification &spec); + /// + /// \param spec[in] Task specification. + /// \param results[out] Results for task execution. + /// \return Status. + Status ExecuteTask(const TaskSpecification &spec, + std::vector> *results); /// Reference to the parent CoreWorker's context. WorkerContext &worker_context_; @@ -66,7 +73,8 @@ class CoreWorkerTaskExecutionInterface { TaskExecutor execution_callback_; /// All the task task receivers supported. - std::unordered_map> task_receivers_; + EnumUnorderedMap> + task_receivers_; /// The RPC server. rpc::GrpcServer worker_server_; diff --git a/src/ray/core_worker/task_interface.cc b/src/ray/core_worker/task_interface.cc index c40b2c05d..47345d79b 100644 --- a/src/ray/core_worker/task_interface.cc +++ b/src/ray/core_worker/task_interface.cc @@ -93,7 +93,7 @@ void ActorHandle::ClearNewActorHandles() { new_actor_handles_.clear(); } CoreWorkerTaskInterface::CoreWorkerTaskInterface( WorkerContext &worker_context, std::unique_ptr &raylet_client) : worker_context_(worker_context) { - task_submitters_.emplace(static_cast(TaskTransportType::RAYLET), + task_submitters_.emplace(TaskTransportType::RAYLET, std::unique_ptr( new CoreWorkerRayletTaskSubmitter(raylet_client))); } @@ -134,8 +134,7 @@ Status CoreWorkerTaskInterface::SubmitTask(const RayFunction &function, std::vector *return_ids) { auto builder = BuildCommonTaskSpec(function, args, task_options.num_returns, task_options.resources, {}, return_ids); - return task_submitters_[static_cast(TaskTransportType::RAYLET)]->SubmitTask( - builder.Build()); + return task_submitters_[TaskTransportType::RAYLET]->SubmitTask(builder.Build()); } Status CoreWorkerTaskInterface::CreateActor( @@ -155,8 +154,7 @@ Status CoreWorkerTaskInterface::CreateActor( (*actor_handle)->IncreaseTaskCounter(); (*actor_handle)->SetActorCursor(return_ids[0]); - return task_submitters_[static_cast(TaskTransportType::RAYLET)]->SubmitTask( - builder.Build()); + return task_submitters_[TaskTransportType::RAYLET]->SubmitTask(builder.Build()); } Status CoreWorkerTaskInterface::SubmitActorTask(ActorHandle &actor_handle, @@ -188,8 +186,7 @@ Status CoreWorkerTaskInterface::SubmitActorTask(ActorHandle &actor_handle, guard.unlock(); // Submit task. - auto status = task_submitters_[static_cast(TaskTransportType::RAYLET)]->SubmitTask( - builder.Build()); + auto status = task_submitters_[TaskTransportType::RAYLET]->SubmitTask(builder.Build()); // Remove cursor from return ids. (*return_ids).pop_back(); diff --git a/src/ray/core_worker/task_interface.h b/src/ray/core_worker/task_interface.h index 2fea4fd68..e1ec502b6 100644 --- a/src/ray/core_worker/task_interface.h +++ b/src/ray/core_worker/task_interface.h @@ -171,7 +171,8 @@ class CoreWorkerTaskInterface { WorkerContext &worker_context_; /// All the task submitters supported. - std::unordered_map> task_submitters_; + EnumUnorderedMap> + task_submitters_; }; } // namespace ray diff --git a/src/ray/core_worker/transport/raylet_transport.cc b/src/ray/core_worker/transport/raylet_transport.cc index 4698d4363..9706b2ce4 100644 --- a/src/ray/core_worker/transport/raylet_transport.cc +++ b/src/ray/core_worker/transport/raylet_transport.cc @@ -14,9 +14,11 @@ Status CoreWorkerRayletTaskSubmitter::SubmitTask(const TaskSpecification &task) } CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver( - std::unique_ptr &raylet_client, boost::asio::io_service &io_service, + std::unique_ptr &raylet_client, + CoreWorkerObjectInterface &object_interface, boost::asio::io_service &io_service, rpc::GrpcServer &server, const TaskHandler &task_handler) : raylet_client_(raylet_client), + object_interface_(object_interface), task_service_(io_service, *this), task_handler_(task_handler) { server.RegisterService(task_service_); @@ -27,7 +29,22 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask( rpc::SendReplyCallback send_reply_callback) { const Task task(request.task()); const auto &task_spec = task.GetTaskSpecification(); - auto status = task_handler_(task_spec); + std::vector> results; + auto status = task_handler_(task_spec, &results); + + auto num_returns = task_spec.NumReturns(); + if (task_spec.IsActorCreationTask() || task_spec.IsActorTask()) { + RAY_CHECK(num_returns > 0); + // Decrease to account for the dummy object id. + num_returns--; + } + + RAY_CHECK(results.size() == num_returns); + for (int i = 0; i < num_returns; i++) { + ObjectID id = ObjectID::ForTaskReturn(task_spec.TaskId(), i + 1); + object_interface_.Put(*results[i], id); + } + // Notify raylet that current task is done via a `TaskDone` message. This is to // ensure that the task is marked as finished by raylet only after previous // raylet client calls are completed. For example, if the worker sends a diff --git a/src/ray/core_worker/transport/raylet_transport.h b/src/ray/core_worker/transport/raylet_transport.h index 72b0f5ffc..0ba8feb5e 100644 --- a/src/ray/core_worker/transport/raylet_transport.h +++ b/src/ray/core_worker/transport/raylet_transport.h @@ -3,6 +3,7 @@ #include +#include "ray/core_worker/object_interface.h" #include "ray/core_worker/transport/transport.h" #include "ray/raylet/raylet_client.h" #include "ray/rpc/worker/worker_server.h" @@ -32,6 +33,7 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver, public rpc::WorkerTaskHandler { public: CoreWorkerRayletTaskReceiver(std::unique_ptr &raylet_client, + CoreWorkerObjectInterface &object_interface, boost::asio::io_service &io_service, rpc::GrpcServer &server, const TaskHandler &task_handler); @@ -49,6 +51,8 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver, private: /// Raylet client. std::unique_ptr &raylet_client_; + // Object interface. + CoreWorkerObjectInterface &object_interface_; /// The rpc service for `WorkerTaskService`. rpc::WorkerTaskGrpcService task_service_; /// The callback function to process a task. diff --git a/src/ray/core_worker/transport/transport.h b/src/ray/core_worker/transport/transport.h index e103fa130..4fd3e7763 100644 --- a/src/ray/core_worker/transport/transport.h +++ b/src/ray/core_worker/transport/transport.h @@ -8,6 +8,7 @@ #include "ray/common/status.h" #include "ray/common/task/task_spec.h" #include "ray/core_worker/common.h" +#include "ray/core_worker/store_provider/store_provider.h" namespace ray { @@ -32,7 +33,9 @@ class CoreWorkerTaskSubmitter { /// This class receives tasks for execution. class CoreWorkerTaskReceiver { public: - using TaskHandler = std::function; + using TaskHandler = + std::function> *results)>; }; } // namespace ray diff --git a/src/ray/util/util.h b/src/ray/util/util.h index f86870858..c8c9930ff 100644 --- a/src/ray/util/util.h +++ b/src/ray/util/util.h @@ -3,6 +3,7 @@ #include #include +#include #include "ray/common/status.h" @@ -78,4 +79,15 @@ class InitShutdownRAII { ShutdownFunc shutdown_; }; +struct EnumClassHash { + template + std::size_t operator()(T t) const { + return static_cast(t); + } +}; + +/// unodered_map for enum class type. +template +using EnumUnorderedMap = std::unordered_map; + #endif // RAY_UTIL_UTIL_H