[direct call] changes raylet to push tasks to worker

src/ray/core_worker/core_worker.cc

@@ -6,7 +6,7 @@ namespace ray {
 CoreWorker::CoreWorker(const enum WorkerType worker_type, const enum Language language,
                        const std::string &store_socket, const std::string &raylet_socket,
                        const JobID &job_id,
-                       CoreWorkerTaskExecutionInterface::TaskExecutor execution_callback)
+                       const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback)
     : worker_type_(worker_type),
       language_(language),
       raylet_socket_(raylet_socket),

src/ray/core_worker/core_worker.h

@@ -25,7 +25,7 @@ class CoreWorker {
   CoreWorker(const WorkerType worker_type, const Language language,
              const std::string &store_socket, const std::string &raylet_socket,
              const JobID &job_id,
-             CoreWorkerTaskExecutionInterface::TaskExecutor execution_callback = nullptr);
+             const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback);
 
   /// Type of this worker.
   enum WorkerType WorkerType() const { return worker_type_; }

src/ray/core_worker/core_worker_test.cc

@@ -125,7 +125,7 @@ class CoreWorkerTest : public ::testing::Test {
 
   void TestNormalTask(const std::unordered_map<std::string, double> &resources) {
     CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
-                      raylet_socket_names_[0], JobID::FromRandom());
+                      raylet_socket_names_[0], JobID::FromRandom(), nullptr);
 
     // Test pass by value.
     {
@@ -184,7 +184,7 @@ class CoreWorkerTest : public ::testing::Test {
 
   void TestActorTask(const std::unordered_map<std::string, double> &resources) {
     CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
-                      raylet_socket_names_[0], JobID::FromRandom());
+                      raylet_socket_names_[0], JobID::FromRandom(), nullptr);
 
     std::unique_ptr<ActorHandle> actor_handle;
 
@@ -335,7 +335,7 @@ TEST_F(ZeroNodeTest, TestActorHandle) {
 TEST_F(SingleNodeTest, TestObjectInterface) {
   CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON,
                          raylet_store_socket_names_[0], raylet_socket_names_[0],
-                         JobID::FromRandom());
+                         JobID::FromRandom(), nullptr);
 
   uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8};
   uint8_t array2[] = {10, 11, 12, 13, 14, 15};
@@ -398,10 +398,10 @@ TEST_F(SingleNodeTest, TestObjectInterface) {
 
 TEST_F(TwoNodeTest, TestObjectInterfaceCrossNodes) {
   CoreWorker worker1(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
-                     raylet_socket_names_[0], JobID::FromRandom());
+                     raylet_socket_names_[0], JobID::FromRandom(), nullptr);
 
   CoreWorker worker2(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[1],
-                     raylet_socket_names_[1], JobID::FromRandom());
+                     raylet_socket_names_[1], JobID::FromRandom(), nullptr);
 
   uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8};
   uint8_t array2[] = {10, 11, 12, 13, 14, 15};
@@ -487,7 +487,7 @@ TEST_F(TwoNodeTest, TestActorTaskCrossNodes) {
 TEST_F(SingleNodeTest, TestCoreWorkerConstructorFailure) {
   try {
     CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON, "",
-                           raylet_socket_names_[0], JobID::FromRandom());
+                           raylet_socket_names_[0], JobID::FromRandom(), nullptr);
   } catch (const std::exception &e) {
     std::cout << "Caught exception when constructing core worker: " << e.what();
   }

src/ray/core_worker/mock_worker.cc

@@ -3,6 +3,8 @@
 #include "ray/core_worker/store_provider/store_provider.h"
 #include "ray/core_worker/task_execution.h"
 
+using namespace std::placeholders;
+
 namespace ray {
 
 /// A mock C++ worker used by core_worker_test.cc to verify the task submission/execution
@@ -16,43 +18,42 @@ namespace ray {
 /// for more details on how this class is used.
 class MockWorker {
  public:
-  MockWorker(const std::string &store_socket, const std::string &raylet_socket) {
-    auto executor_func = [this](const RayFunction &ray_function,
-                                const std::vector<std::shared_ptr<RayObject>> &args,
-                                const TaskInfo &task_info, int num_returns) {
-      // Note that this doesn't include dummy object id.
-      RAY_CHECK(num_returns >= 0);
+  MockWorker(const std::string &store_socket, const std::string &raylet_socket)
+      : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
+                JobID::FromRandom(),
+                std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4)) {}
 
-      // Merge all the content from input args.
-      std::vector<uint8_t> buffer;
-      for (const auto &arg : args) {
-        auto &data = arg->GetData();
-        buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size());
-      }
+  void Run() {
+    // Start executing tasks.
+    worker_.Execution().Run();
+  }
 
-      auto return_value = RayObject(
-          std::make_shared<LocalMemoryBuffer>(buffer.data(), buffer.size()), nullptr);
+ private:
+  Status ExecuteTask(const RayFunction &ray_function,
+                     const std::vector<std::shared_ptr<RayObject>> &args,
+                     const TaskInfo &task_info, int num_returns) {
+    // Note that this doesn't include dummy object id.
+    RAY_CHECK(num_returns >= 0);
 
-      // 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));
-      }
-      return Status::OK();
-    };
+    // Merge all the content from input args.
+    std::vector<uint8_t> buffer;
+    for (const auto &arg : args) {
+      auto &data = arg->GetData();
+      buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size());
+    }
 
-    worker_ = std::unique_ptr<CoreWorker>(
-        new CoreWorker(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
-                       JobID::FromRandom(), executor_func));
-  }
+    auto return_value = RayObject(
+        std::make_shared<LocalMemoryBuffer>(buffer.data(), buffer.size()), nullptr);
 
-  void Run() {
-    // Start executing tasks.
-    worker_->Execution().Run();
+    // 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));
+    }
+    return Status::OK();
   }
 
- private:
-  std::unique_ptr<CoreWorker> worker_;
+  CoreWorker worker_;
 };
 
 }  // namespace ray

src/ray/core_worker/task_execution.cc

@@ -19,8 +19,8 @@ CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface(
                         std::placeholders::_1);
   task_receivers_.emplace(
       static_cast<int>(TaskTransportType::RAYLET),
-      std::unique_ptr<CoreWorkerRayletTaskReceiver>(
-          new CoreWorkerRayletTaskReceiver(main_service_, worker_server_, func)));
+      std::unique_ptr<CoreWorkerRayletTaskReceiver>(new CoreWorkerRayletTaskReceiver(
+          raylet_client, main_service_, worker_server_, func)));
 
   // Start RPC server after all the task receivers are properly initialized.
   worker_server_.Run();
@@ -60,12 +60,12 @@ Status CoreWorkerTaskExecutionInterface::ExecuteTask(
   return status;
 }
 
-Status CoreWorkerTaskExecutionInterface::Run() {
+void CoreWorkerTaskExecutionInterface::Run() {
   // Run main IO service.
   main_service_.run();
 
   // should never reach here.
-  return Status::OK();
+  RAY_LOG(FATAL) << "should never reach here after running main io service";
 }
 
 Status CoreWorkerTaskExecutionInterface::BuildArgsForExecutor(

src/ray/core_worker/task_execution.h

@@ -39,8 +39,8 @@ class CoreWorkerTaskExecutionInterface {
                                    const TaskExecutor &executor);
 
   /// Start receving and executes tasks in a infinite loop.
-  /// \return Status.
-  Status Run();
+  /// \return void.
+  void Run();
 
  private:
   /// Build arguments for task executor. This would loop through all the arguments

src/ray/core_worker/transport/raylet_transport.cc

@@ -14,9 +14,11 @@ Status CoreWorkerRayletTaskSubmitter::SubmitTask(const TaskSpec &task) {
 }
 
 CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver(
-    boost::asio::io_service &io_service, rpc::GrpcServer &server,
-    const TaskHandler &task_handler)
-    : task_service_(io_service, *this), task_handler_(task_handler) {
+    std::unique_ptr<RayletClient> &raylet_client, boost::asio::io_service &io_service,
+    rpc::GrpcServer &server, const TaskHandler &task_handler)
+    : raylet_client_(raylet_client),
+      task_service_(io_service, *this),
+      task_handler_(task_handler) {
   server.RegisterService(task_service_);
 }
 
@@ -26,6 +28,12 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
   const raylet::Task task(request.task());
   const auto &spec = task.GetTaskSpecification();
   auto status = task_handler_(spec);
+  // Notify raylet the current task is done. This is to ensure that the task
+  // is marked as finished by raylet only after previous raylet client calls are
+  // completed. The rpc `done_callback` is sent via a different connection
+  // from raylet client connection, so it cannot guarantee the rpc reply arrives
+  // at raylet after a previous `NotifyUnblocked` message.
+  raylet_client_->TaskDone();
   send_reply_callback(status, nullptr, nullptr);
 }
 

src/ray/core_worker/transport/raylet_transport.h

@@ -31,7 +31,8 @@ class CoreWorkerRayletTaskSubmitter : public CoreWorkerTaskSubmitter {
 class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver,
                                      public rpc::WorkerTaskHandler {
  public:
-  CoreWorkerRayletTaskReceiver(boost::asio::io_service &io_service,
+  CoreWorkerRayletTaskReceiver(std::unique_ptr<RayletClient> &raylet_client,
+                               boost::asio::io_service &io_service,
                                rpc::GrpcServer &server, const TaskHandler &task_handler);
 
   /// Handle a `AssignTask` request.
@@ -46,6 +47,8 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver,
                         rpc::SendReplyCallback send_reply_callback) override;
 
  private:
+  /// Raylet client.
+  std::unique_ptr<RayletClient> &raylet_client_;
   /// The rpc service for `WorkerTaskService`.
   rpc::WorkerTaskGrpcService task_service_;
   /// The callback function to process a task.

src/ray/protobuf/worker.proto

@@ -5,18 +5,12 @@ package ray.rpc;
 import "src/ray/protobuf/common.proto";
 
 message AssignTaskRequest {
-  // The ID of the task to be pushed.
-  bytes task_id = 1;
-  // The task to be assigned. This should include task_id.
-  // TODO(hchen): Currently, `task_spec` are represented as
-  // flatbutters-serialized bytes. This is because the flatbuffers-defined Task data
-  // structure is being used in many places. We should move Task and all related data
-  // structures to protobuf.
-  bytes task_spec = 2;
+  // The task to be pushed.
+  Task task = 1;
   // A list of the resources reserved for this worker.
   // TODO(zhijunfu): `resource_ids` is represented as
   // flatbutters-serialized bytes, will be moved to protobuf later.
-  bytes resource_ids = 3;
+  bytes resource_ids = 2;
 }
 
 message AssignTaskReply {