[direct call] add local plasma provider (#5184)

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<uint8_t> buffer_;
 };
 
 /// Represents a byte buffer for plasma object.

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 };
 

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<std::shared_ptr<RayObject>> &args,
-                     const TaskInfo &task_info, int num_returns) {
+                     const TaskInfo &task_info, int num_returns,
+                     std::vector<std::shared_ptr<RayObject>> *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<LocalMemoryBuffer>(buffer.data(), buffer.size()), nullptr);
+    auto memory_buffer =
+        std::make_shared<LocalMemoryBuffer>(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<RayObject>(memory_buffer, nullptr));
     }
     return Status::OK();
   }

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<RayletClient> &raylet_client,
     const std::string &store_socket)
-    : worker_context_(worker_context), raylet_client_(raylet_client) {
-  store_providers_.emplace(
-      static_cast<int>(StoreProviderType::PLASMA),
-      std::unique_ptr<CoreWorkerStoreProvider>(
-          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<int>(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<ObjectID> &ids,
                                       int64_t timeout_ms,
                                       std::vector<std::shared_ptr<RayObject>> *results) {
-  auto type = static_cast<int>(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<ObjectID> &object_ids,
                                        int num_objects, int64_t timeout_ms,
                                        std::vector<bool> *results) {
-  auto type = static_cast<int>(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<ObjectID> &object_ids,
                                          bool local_only, bool delete_creating_tasks) {
-  auto type = static_cast<int>(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<CoreWorkerStoreProvider> CoreWorkerObjectInterface::CreateStoreProvider(
+    StoreProviderType type) const {
+  switch (type) {
+  case StoreProviderType::LOCAL_PLASMA:
+    return std::unique_ptr<CoreWorkerStoreProvider>(
+        new CoreWorkerLocalPlasmaStoreProvider(store_socket_));
+    break;
+  case StoreProviderType::PLASMA:
+    return std::unique_ptr<CoreWorkerStoreProvider>(
+        new CoreWorkerPlasmaStoreProvider(store_socket_, raylet_client_));
+    break;
+  default:
+    RAY_LOG(FATAL) << "unknown store provider type " << static_cast<int>(type);
+    break;
+  }
 }
 
 }  // namespace ray

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<CoreWorkerStoreProvider> 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<RayletClient> &raylet_client_;
 
+  /// Store socket name.
+  std::string store_socket_;
+
   /// All the store providers supported.
-  std::unordered_map<int, std::unique_ptr<CoreWorkerStoreProvider>> store_providers_;
+  EnumUnorderedMap<StoreProviderType, std::unique_ptr<CoreWorkerStoreProvider>>
+      store_providers_;
+
+  friend class CoreWorkerTaskInterface;
 };
 
 }  // namespace ray

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<arrow::Buffer> out_buffer;
+  {
+    std::unique_lock<std::mutex> 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<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 != 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]->GetData() != 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

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<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

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<RayletClient> &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<arrow::Buffer> out_buffer;
-  {
-    std::unique_lock<std::mutex> 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<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();
+  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::ObjectID> plasma_ids;
-    for (const auto &id : unready_ids) {
-      plasma_ids.push_back(id.ToPlasmaId());
-    }
+    std::vector<std::shared_ptr<RayObject>> result_objects;
+    RAY_RETURN_NOT_OK(
+        local_store_provider_.Get(unready_ids, get_timeout, task_id, &result_objects));
 
-    std::vector<plasma::ObjectBuffer> object_buffers;
-    {
-      std::unique_lock<std::mutex> 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<RayObject>(
-            std::make_shared<PlasmaBuffer>(object_buffers[i].data),
-            std::make_shared<PlasmaBuffer>(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<ObjectID> &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<const char *>(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();

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<ObjectID> &object_ids, bool local_only,
-                bool delete_creating_tasks) override;
+  /// created these objects.
+  /// \return Status.
+  Status Delete(const std::vector<ObjectID> &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<ObjectID, int> &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<RayletClient> &raylet_client_;

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<ObjectID> &object_ids, bool local_only,
-                        bool delete_creating_tasks) = 0;
+  virtual Status Delete(const std::vector<ObjectID> &object_ids, bool local_only = true,
+                        bool delete_creating_tasks = false) = 0;
 };
 
 }  // namespace ray

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<int>(TaskTransportType::RAYLET),
+      TaskTransportType::RAYLET,
       std::unique_ptr<CoreWorkerRayletTaskReceiver>(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<std::shared_ptr<RayObject>> *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.

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<Status(const RayFunction &ray_function,
-                           const std::vector<std::shared_ptr<RayObject>> &args,
-                           const TaskInfo &task_info, int num_returns)>;
+  using TaskExecutor = std::function<Status(
+      const RayFunction &ray_function,
+      const std::vector<std::shared_ptr<RayObject>> &args, const TaskInfo &task_info,
+      int num_returns, std::vector<std::shared_ptr<RayObject>> *results)>;
 
   CoreWorkerTaskExecutionInterface(WorkerContext &worker_context,
                                    std::unique_ptr<RayletClient> &raylet_client,
@@ -55,7 +57,12 @@ class CoreWorkerTaskExecutionInterface {
                               std::vector<std::shared_ptr<RayObject>> *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<std::shared_ptr<RayObject>> *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<int, std::unique_ptr<CoreWorkerTaskReceiver>> task_receivers_;
+  EnumUnorderedMap<TaskTransportType, std::unique_ptr<CoreWorkerTaskReceiver>>
+      task_receivers_;
 
   /// The RPC server.
   rpc::GrpcServer worker_server_;

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<RayletClient> &raylet_client)
     : worker_context_(worker_context) {
-  task_submitters_.emplace(static_cast<int>(TaskTransportType::RAYLET),
+  task_submitters_.emplace(TaskTransportType::RAYLET,
                            std::unique_ptr<CoreWorkerRayletTaskSubmitter>(
                                new CoreWorkerRayletTaskSubmitter(raylet_client)));
 }
@@ -134,8 +134,7 @@ Status CoreWorkerTaskInterface::SubmitTask(const RayFunction &function,
                                            std::vector<ObjectID> *return_ids) {
   auto builder = BuildCommonTaskSpec(function, args, task_options.num_returns,
                                      task_options.resources, {}, return_ids);
-  return task_submitters_[static_cast<int>(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<int>(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<int>(TaskTransportType::RAYLET)]->SubmitTask(
-      builder.Build());
+  auto status = task_submitters_[TaskTransportType::RAYLET]->SubmitTask(builder.Build());
 
   // Remove cursor from return ids.
   (*return_ids).pop_back();

src/ray/core_worker/task_interface.h

@@ -171,7 +171,8 @@ class CoreWorkerTaskInterface {
   WorkerContext &worker_context_;
 
   /// All the task submitters supported.
-  std::unordered_map<int, std::unique_ptr<CoreWorkerTaskSubmitter>> task_submitters_;
+  EnumUnorderedMap<TaskTransportType, std::unique_ptr<CoreWorkerTaskSubmitter>>
+      task_submitters_;
 };
 
 }  // namespace ray

src/ray/core_worker/transport/raylet_transport.cc

@@ -14,9 +14,11 @@ Status CoreWorkerRayletTaskSubmitter::SubmitTask(const TaskSpecification &task)
 }
 
 CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver(
-    std::unique_ptr<RayletClient> &raylet_client, boost::asio::io_service &io_service,
+    std::unique_ptr<RayletClient> &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<std::shared_ptr<RayObject>> 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

src/ray/core_worker/transport/raylet_transport.h

@@ -3,6 +3,7 @@
 
 #include <list>
 
+#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<RayletClient> &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<RayletClient> &raylet_client_;
+  // Object interface.
+  CoreWorkerObjectInterface &object_interface_;
   /// The rpc service for `WorkerTaskService`.
   rpc::WorkerTaskGrpcService task_service_;
   /// The callback function to process a task.

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<Status(const TaskSpecification &task_spec)>;
+  using TaskHandler =
+      std::function<Status(const TaskSpecification &task_spec,
+                           std::vector<std::shared_ptr<RayObject>> *results)>;
 };
 
 }  // namespace ray

src/ray/util/util.h

@@ -3,6 +3,7 @@
 
 #include <boost/system/error_code.hpp>
 #include <chrono>
+#include <unordered_map>
 
 #include "ray/common/status.h"
 
@@ -78,4 +79,15 @@ class InitShutdownRAII {
   ShutdownFunc shutdown_;
 };
 
+struct EnumClassHash {
+  template <typename T>
+  std::size_t operator()(T t) const {
+    return static_cast<std::size_t>(t);
+  }
+};
+
+/// unodered_map for enum class type.
+template <typename Key, typename T>
+using EnumUnorderedMap = std::unordered_map<Key, T, EnumClassHash>;
+
 #endif  // RAY_UTIL_UTIL_H