[direct call] changes raylet to push tasks to worker (#5140)

* refactor grpc server

* format

* change GetTask() to PushTask()

* change PushTask to AssignTask

* format

* add resource_ids

* move done_callback to server call

* remove SetTaskHandler and initialize it in task receiver's constructor

* format

* resolve comments

* update

* update

* Update src/ray/core_worker/core_worker.cc

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* resolve comments

* format

* Update src/ray/core_worker/transport/raylet_transport.cc

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* resolve comments

* resolve comments

* fix build

* format

* fix

* format

* noop

src/ray/common/task/scheduling_resources.cc

@@ -700,6 +700,13 @@ std::vector<flatbuffers::Offset<protocol::ResourceIdSetInfo>> ResourceIdSet::ToF
   return return_message;
 }
 
+const std::string ResourceIdSet::Serialize() const {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto resource_id_set_flatbuf = ToFlatbuf(fbb);
+  fbb.Finish(fbb.CreateVector(resource_id_set_flatbuf));
+  return std::string(fbb.GetBufferPointer(), fbb.GetBufferPointer() + fbb.GetSize());
+}
+
 /// SchedulingResources class implementation
 
 SchedulingResources::SchedulingResources()

src/ray/common/task/scheduling_resources.h

@@ -429,6 +429,12 @@ class ResourceIdSet {
   std::vector<flatbuffers::Offset<ray::protocol::ResourceIdSetInfo>> ToFlatbuf(
       flatbuffers::FlatBufferBuilder &fbb) const;
 
+  /// \brief Serialize this object as a string.
+  ///
+  /// \return A serialized string of this object.
+  /// TODO(zhijunfu): this can be removed after raylet client is migrated to grpc.
+  const std::string Serialize() const;
+
  private:
   /// A mapping from reosurce name to a set of resource IDs for that resource.
   std::unordered_map<std::string, ResourceIds> available_resources_;

src/ray/core_worker/core_worker.cc

@@ -3,9 +3,11 @@
 
 namespace ray {
 
-CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
-                       const std::string &store_socket, const std::string &raylet_socket,
-                       const JobID &job_id)
+CoreWorker::CoreWorker(
+    const WorkerType worker_type, const Language language,
+    const std::string &store_socket, const std::string &raylet_socket,
+    const JobID &job_id,
+    const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback)
     : worker_type_(worker_type),
       language_(language),
       raylet_socket_(raylet_socket),
@@ -14,9 +16,10 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
       object_interface_(worker_context_, raylet_client_, store_socket) {
   int rpc_server_port = 0;
   if (worker_type_ == ray::WorkerType::WORKER) {
+    RAY_CHECK(execution_callback != nullptr);
     task_execution_interface_ = std::unique_ptr<CoreWorkerTaskExecutionInterface>(
         new CoreWorkerTaskExecutionInterface(worker_context_, raylet_client_,
-                                             object_interface_));
+                                             object_interface_, execution_callback));
     rpc_server_port = task_execution_interface_->worker_server_.GetPort();
   }
   // TODO(zhijunfu): currently RayletClient would crash in its constructor if it cannot

src/ray/core_worker/core_worker.h

@@ -24,7 +24,8 @@ class CoreWorker {
   /// NOTE(zhijunfu): the constructor would throw if a failure happens.
   CoreWorker(const WorkerType worker_type, const Language language,
              const std::string &store_socket, const std::string &raylet_socket,
-             const JobID &job_id = JobID::Nil());
+             const JobID &job_id,
+             const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback);
 
   /// Type of this worker.
   WorkerType GetWorkerType() 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::FromInt(1));
+                      raylet_socket_names_[0], JobID::FromInt(1), 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::FromInt(1));
+                      raylet_socket_names_[0], JobID::FromInt(1), 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::JobID::FromInt(1));
+                         JobID::FromInt(1), 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::JobID::FromInt(1));
+                     raylet_socket_names_[0], JobID::FromInt(1), nullptr);
 
   CoreWorker worker2(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[1],
-                     raylet_socket_names_[1], JobID::JobID::FromInt(1));
+                     raylet_socket_names_[1], JobID::FromInt(1), 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::FromInt(1));
+                           raylet_socket_names_[0], JobID::FromInt(1), 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
@@ -18,38 +20,39 @@ class MockWorker {
  public:
   MockWorker(const std::string &store_socket, const std::string &raylet_socket)
       : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
-                JobID::JobID::FromInt(1)) {}
+                JobID::JobID::FromInt(1),
+                std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4)) {}
 
   void Run() {
-    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);
-
-      // 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());
-      }
-
-      auto return_value = RayObject(
-          std::make_shared<LocalMemoryBuffer>(buffer.data(), buffer.size()), nullptr);
-
-      // 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();
-    };
-
     // Start executing tasks.
-    worker_.Execution().Run(executor_func);
+    worker_.Execution().Run();
   }
 
  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);
+
+    // 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());
+    }
+
+    auto return_value = RayObject(
+        std::make_shared<LocalMemoryBuffer>(buffer.data(), buffer.size()), nullptr);
+
+    // 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();
+  }
+
   CoreWorker worker_;
 };
 

src/ray/core_worker/task_execution.cc

@@ -7,67 +7,64 @@ namespace ray {
 
 CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface(
     WorkerContext &worker_context, std::unique_ptr<RayletClient> &raylet_client,
-    CoreWorkerObjectInterface &object_interface)
+    CoreWorkerObjectInterface &object_interface, const TaskExecutor &executor)
     : worker_context_(worker_context),
       object_interface_(object_interface),
+      execution_callback_(executor),
       worker_server_("Worker", 0 /* let grpc choose port */),
       main_work_(main_service_) {
+  RAY_CHECK(execution_callback_ != nullptr);
+
+  auto func = std::bind(&CoreWorkerTaskExecutionInterface::ExecuteTask, this,
+                        std::placeholders::_1);
   task_receivers_.emplace(
       static_cast<int>(TaskTransportType::RAYLET),
       std::unique_ptr<CoreWorkerRayletTaskReceiver>(new CoreWorkerRayletTaskReceiver(
-          raylet_client, main_service_, worker_server_)));
+          raylet_client, main_service_, worker_server_, func)));
 
   // Start RPC server after all the task receivers are properly initialized.
   worker_server_.Run();
 }
 
-Status CoreWorkerTaskExecutionInterface::Run(const TaskExecutor &executor) {
-  while (true) {
-    std::vector<TaskSpec> tasks;
-    auto status =
-        task_receivers_[static_cast<int>(TaskTransportType::RAYLET)]->GetTasks(&tasks);
-    if (!status.ok()) {
-      RAY_LOG(ERROR) << "Getting task failed with error: "
-                     << ray::Status::IOError(status.message());
-      return status;
-    }
+Status CoreWorkerTaskExecutionInterface::ExecuteTask(const TaskSpecification &spec) {
+  worker_context_.SetCurrentTask(spec);
 
-    for (const auto &task : tasks) {
-      const auto &spec = task.GetTaskSpecification();
-      worker_context_.SetCurrentTask(spec);
+  RayFunction func{spec.GetLanguage(), spec.FunctionDescriptor()};
 
-      RayFunction func{spec.GetLanguage(), spec.FunctionDescriptor()};
+  std::vector<std::shared_ptr<RayObject>> args;
+  RAY_CHECK_OK(BuildArgsForExecutor(spec, &args));
 
-      std::vector<std::shared_ptr<RayObject>> args;
-      RAY_CHECK_OK(BuildArgsForExecutor(spec, &args));
-
-      TaskType task_type;
-      if (spec.IsActorCreationTask()) {
-        task_type = TaskType::ACTOR_CREATION_TASK;
-      } else if (spec.IsActorTask()) {
-        task_type = TaskType::ACTOR_TASK;
-      } else {
-        task_type = TaskType::NORMAL_TASK;
-      }
-
-      TaskInfo task_info{spec.TaskId(), spec.JobId(), task_type};
-
-      auto num_returns = spec.NumReturns();
-      if (spec.IsActorCreationTask() || spec.IsActorTask()) {
-        RAY_CHECK(num_returns > 0);
-        // Decrease to account for the dummy object id.
-        num_returns--;
-      }
-
-      status = executor(func, args, task_info, num_returns);
-      // TODO(zhijunfu):
-      // 1. Check and handle failure.
-      // 2. Save or load checkpoint.
-    }
+  TaskType task_type;
+  if (spec.IsActorCreationTask()) {
+    task_type = TaskType::ACTOR_CREATION_TASK;
+  } else if (spec.IsActorTask()) {
+    task_type = TaskType::ACTOR_TASK;
+  } else {
+    task_type = TaskType::NORMAL_TASK;
   }
 
+  TaskInfo task_info{spec.TaskId(), spec.JobId(), task_type};
+
+  auto num_returns = spec.NumReturns();
+  if (spec.IsActorCreationTask() || spec.IsActorTask()) {
+    RAY_CHECK(num_returns > 0);
+    // Decrease to account for the dummy object id.
+    num_returns--;
+  }
+
+  auto status = execution_callback_(func, args, task_info, num_returns);
+  // TODO(zhijunfu):
+  // 1. Check and handle failure.
+  // 2. Save or load checkpoint.
+  return status;
+}
+
+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

@@ -23,10 +23,6 @@ class TaskSpecification;
 /// execution.
 class CoreWorkerTaskExecutionInterface {
  public:
-  CoreWorkerTaskExecutionInterface(WorkerContext &worker_context,
-                                   std::unique_ptr<RayletClient> &raylet_client,
-                                   CoreWorkerObjectInterface &object_interface);
-
   /// The callback provided app-language workers that executes tasks.
   ///
   /// \param ray_function[in] Information about the function to execute.
@@ -37,9 +33,14 @@ class CoreWorkerTaskExecutionInterface {
                            const std::vector<std::shared_ptr<RayObject>> &args,
                            const TaskInfo &task_info, int num_returns)>;
 
+  CoreWorkerTaskExecutionInterface(WorkerContext &worker_context,
+                                   std::unique_ptr<RayletClient> &raylet_client,
+                                   CoreWorkerObjectInterface &object_interface,
+                                   const TaskExecutor &executor);
+
   /// Start receving and executes tasks in a infinite loop.
-  /// \return Status.
-  Status Run(const TaskExecutor &executor);
+  /// \return void.
+  void Run();
 
  private:
   /// Build arguments for task executor. This would loop through all the arguments
@@ -53,11 +54,17 @@ class CoreWorkerTaskExecutionInterface {
   Status BuildArgsForExecutor(const TaskSpecification &spec,
                               std::vector<std::shared_ptr<RayObject>> *args);
 
+  /// Execute a task.
+  Status ExecuteTask(const TaskSpecification &spec);
+
   /// Reference to the parent CoreWorker's context.
   WorkerContext &worker_context_;
   /// Reference to the parent CoreWorker's objects interface.
   CoreWorkerObjectInterface &object_interface_;
 
+  // Task execution callback.
+  TaskExecutor execution_callback_;
+
   /// All the task task receivers supported.
   std::unordered_map<int, std::unique_ptr<CoreWorkerTaskReceiver>> task_receivers_;
 

src/ray/core_worker/transport/raylet_transport.cc

@@ -9,29 +9,16 @@ CoreWorkerRayletTaskSubmitter::CoreWorkerRayletTaskSubmitter(
     : raylet_client_(raylet_client) {}
 
 Status CoreWorkerRayletTaskSubmitter::SubmitTask(const TaskSpec &task) {
+  RAY_CHECK(raylet_client_ != nullptr);
   return raylet_client_->SubmitTask(task.GetDependencies(), task.GetTaskSpecification());
 }
 
-Status CoreWorkerRayletTaskReceiver::GetTasks(std::vector<TaskSpec> *tasks) {
-  std::unique_ptr<TaskSpecification> task_spec;
-  auto status = raylet_client_->GetTask(&task_spec);
-  if (!status.ok()) {
-    RAY_LOG(ERROR) << "Get task from raylet failed with error: "
-                   << ray::Status::IOError(status.message());
-    return status;
-  }
-
-  std::vector<ObjectID> dependencies;
-  RAY_CHECK((*tasks).empty());
-  (*tasks).emplace_back(*task_spec, dependencies);
-
-  return Status::OK();
-}
-
 CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver(
     std::unique_ptr<RayletClient> &raylet_client, boost::asio::io_service &io_service,
-    rpc::GrpcServer &server)
-    : raylet_client_(raylet_client), task_service_(io_service, *this) {
+    rpc::GrpcServer &server, const TaskHandler &task_handler)
+    : raylet_client_(raylet_client),
+      task_service_(io_service, *this),
+      task_handler_(task_handler) {
   server.RegisterService(task_service_);
 }
 
@@ -41,6 +28,19 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
   const Task task(request.task());
   const auto &spec = task.GetTaskSpecification();
   auto status = task_handler_(spec);
+  // 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
+  // NotifyUnblocked message that it is no longer blocked in a `ray.get`
+  // on the normal raylet socket, then completes an assigned task, we
+  // need to guarantee that raylet gets the former message first before
+  // marking the task as completed. This is why a `TaskDone` message
+  // is required - without it, it's possible that raylet receives
+  // rpc reply first before the NotifyUnblocked message arrives,
+  // as they use different connections, the `TaskDone` message is sent
+  // to raylet via the same connection so the order is guaranteed.
+  raylet_client_->TaskDone();
+  // send rpc reply.
   send_reply_callback(status, nullptr, nullptr);
 }
 

src/ray/core_worker/transport/raylet_transport.h

@@ -33,15 +33,8 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver,
  public:
   CoreWorkerRayletTaskReceiver(std::unique_ptr<RayletClient> &raylet_client,
                                boost::asio::io_service &io_service,
-                               rpc::GrpcServer &server);
+                               rpc::GrpcServer &server, const TaskHandler &task_handler);
 
-  // Get tasks for execution from raylet.
-  virtual Status GetTasks(std::vector<TaskSpec> *tasks) override;
-
-  /// TODO(zhijunfu): This is currently unused. Later when we migrate from worker "get
-  /// task" to raylet "assign task", this method will be used and the `GetTask` above will
-  /// be removed.
-  ///
   /// Handle a `AssignTask` request.
   /// The implementation can handle this request asynchronously. When hanling is done, the
   /// `send_reply_callback` should be called.
@@ -56,10 +49,10 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver,
  private:
   /// Raylet client.
   std::unique_ptr<RayletClient> &raylet_client_;
-  /// The callback function to process a task.
-  TaskHandler task_handler_;
   /// The rpc service for `WorkerTaskService`.
   rpc::WorkerTaskGrpcService task_service_;
+  /// The callback function to process a task.
+  TaskHandler task_handler_;
 };
 
 }  // namespace ray

src/ray/core_worker/transport/transport.h

@@ -33,9 +33,6 @@ class CoreWorkerTaskSubmitter {
 class CoreWorkerTaskReceiver {
  public:
   using TaskHandler = std::function<Status(const TaskSpecification &task_spec)>;
-
-  // Get tasks for execution.
-  virtual Status GetTasks(std::vector<TaskSpec> *tasks) = 0;
 };
 
 }  // namespace ray

src/ray/protobuf/worker.proto

@@ -7,6 +7,10 @@ import "src/ray/protobuf/common.proto";
 message AssignTaskRequest {
   // 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 = 2;
 }
 
 message AssignTaskReply {

src/ray/raylet/node_manager.cc

@@ -777,7 +777,12 @@ void NodeManager::ProcessClientMessage(
     ProcessRegisterClientRequestMessage(client, message_data);
   } break;
   case protocol::MessageType::GetTask: {
-    ProcessGetTaskMessage(client);
+    RAY_CHECK(!registered_worker->UsePush());
+    HandleWorkerAvailable(client);
+  } break;
+  case protocol::MessageType::TaskDone: {
+    RAY_CHECK(registered_worker->UsePush());
+    HandleWorkerAvailable(client);
   } break;
   case protocol::MessageType::DisconnectClient: {
     ProcessDisconnectClientMessage(client);
@@ -849,14 +854,20 @@ void NodeManager::ProcessClientMessage(
 void NodeManager::ProcessRegisterClientRequestMessage(
     const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
   auto message = flatbuffers::GetRoot<protocol::RegisterClientRequest>(message_data);
-  client->SetClientID(from_flatbuf<ClientID>(*message->worker_id()));
+  auto client_id = from_flatbuf<ClientID>(*message->worker_id());
+  client->SetClientID(client_id);
   Language language = static_cast<Language>(message->language());
-  auto worker =
-      std::make_shared<Worker>(message->worker_pid(), language, message->port(), client);
+  auto worker = std::make_shared<Worker>(message->worker_pid(), language, message->port(),
+                                         client, client_call_manager_);
   if (message->is_worker()) {
     // Register the new worker.
+    bool use_push_task = worker->UsePush();
+    auto connection = worker->Connection();
     worker_pool_.RegisterWorker(std::move(worker));
-    DispatchTasks(local_queues_.GetReadyTasksWithResources());
+    if (use_push_task) {
+      // only call `HandleWorkerAvailable` when push mode is used.
+      HandleWorkerAvailable(connection);
+    }
   } else {
     // Register the new driver.
     const WorkerID driver_id = from_flatbuf<WorkerID>(*message->worker_id());
@@ -917,7 +928,7 @@ void NodeManager::HandleDisconnectedActor(const ActorID &actor_id, bool was_loca
   PublishActorStateTransition(actor_id, new_actor_data, failure_callback);
 }
 
-void NodeManager::ProcessGetTaskMessage(
+void NodeManager::HandleWorkerAvailable(
     const std::shared_ptr<LocalClientConnection> &client) {
   std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
   RAY_CHECK(worker);
@@ -925,6 +936,7 @@ void NodeManager::ProcessGetTaskMessage(
   if (!worker->GetAssignedTaskId().IsNil()) {
     FinishAssignedTask(*worker);
   }
+
   // Return the worker to the idle pool.
   worker_pool_.PushWorker(std::move(worker));
   // Local resource availability changed: invoke scheduling policy for local node.
@@ -1762,81 +1774,30 @@ bool NodeManager::AssignTask(const Task &task) {
     worker->SetTaskResourceIds(acquired_resources);
   }
 
+  auto task_id = spec.TaskId();
+  auto finish_assign_task_callback = [this, worker, task_id](Status status) {
+    if (worker->UsePush()) {
+      // NOTE: we cannot directly call `FinishAssignTask` here because
+      // it assumes the task is in SWAP queue, thus we need to delay invoking this
+      // function after the assigned tasks are moved from READY queue to SWAP queue
+      // in `DispatchTasks`.
+      // Another option is to move the tasks to SWAP queue here just before calling
+      // `FinishAssignTask` so we can save an io_service post, at the
+      // expense of calling `MoveTask` for each of the assigned tasks.
+      // TODO(zhijunfu): after all workers are fully migrated to push mode, the
+      // `post` below and swap queue can be removed.
+      io_service_.post([this, status, worker, task_id]() {
+        FinishAssignTask(task_id, *worker, status.ok());
+      });
+    } else {
+      FinishAssignTask(task_id, *worker, status.ok());
+    }
+  };
+
   ResourceIdSet resource_id_set =
       worker->GetTaskResourceIds().Plus(worker->GetLifetimeResourceIds());
-  auto resource_id_set_flatbuf = resource_id_set.ToFlatbuf(fbb);
 
-  auto message = protocol::CreateGetTaskReply(fbb, fbb.CreateString(spec.Serialize()),
-                                              fbb.CreateVector(resource_id_set_flatbuf));
-  fbb.Finish(message);
-  const auto &task_id = spec.TaskId();
-  worker->Connection()->WriteMessageAsync(
-      static_cast<int64_t>(protocol::MessageType::ExecuteTask), fbb.GetSize(),
-      fbb.GetBufferPointer(), [this, worker, task_id](ray::Status status) {
-        // Remove the ASSIGNED task from the SWAP queue.
-        Task assigned_task;
-        TaskState state;
-        if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) {
-          return;
-        }
-
-        RAY_CHECK(state == TaskState::SWAP);
-
-        if (status.ok()) {
-          auto spec = assigned_task.GetTaskSpecification();
-          // We successfully assigned the task to the worker.
-          worker->AssignTaskId(spec.TaskId());
-          worker->AssignJobId(spec.JobId());
-          // Actor tasks require extra accounting to track the actor's state.
-          if (spec.IsActorTask()) {
-            auto actor_entry = actor_registry_.find(spec.ActorId());
-            RAY_CHECK(actor_entry != actor_registry_.end());
-            // Process any new actor handles that were created since the
-            // previous task on this handle was executed. The first task
-            // submitted on a new actor handle will depend on the dummy object
-            // returned by the previous task, so the dependency will not be
-            // released until this first task is submitted.
-            for (auto &new_handle_id : spec.NewActorHandles()) {
-              // Get the execution dependency for the first task submitted on the new
-              // actor handle. Since the new actor handle was created after this task
-              // began and before this task finished, it must have the same execution
-              // dependency.
-              const auto &execution_dependencies =
-                  assigned_task.GetTaskExecutionSpec().ExecutionDependencies();
-              // TODO(swang): We expect this task to have exactly 1 execution dependency,
-              // the dummy object returned by the previous actor task. However, this
-              // leaks information about the TaskExecutionSpecification implementation.
-              RAY_CHECK(execution_dependencies.size() == 1);
-              const ObjectID &execution_dependency = execution_dependencies.front();
-              // Add the new handle and give it a reference to the finished task's
-              // execution dependency.
-              actor_entry->second.AddHandle(new_handle_id, execution_dependency);
-            }
-
-            // TODO(swang): For actors with multiple actor handles, to
-            // guarantee that tasks are replayed in the same order after a
-            // failure, we must update the task's execution dependency to be
-            // the actor's current execution dependency.
-          }
-
-          // Mark the task as running.
-          // (See design_docs/task_states.rst for the state transition diagram.)
-          local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING);
-          // Notify the task dependency manager that we no longer need this task's
-          // object dependencies.
-          RAY_CHECK(task_dependency_manager_.UnsubscribeDependencies(spec.TaskId()));
-        } else {
-          RAY_LOG(WARNING) << "Failed to send task to worker, disconnecting client";
-          // We failed to send the task to the worker, so disconnect the worker.
-          ProcessDisconnectClientMessage(worker->Connection());
-          // Queue this task for future assignment. We need to do this since
-          // DispatchTasks() removed it from the ready queue. The task will be
-          // assigned to a worker once one becomes available.
-          // (See design_docs/task_states.rst for the state transition diagram.)
-          local_queues_.QueueTasks({assigned_task}, TaskState::READY);
-          DispatchTasks(MakeTasksWithResources({assigned_task}));
-        }
-      });
+  worker->AssignTask(task, resource_id_set, finish_assign_task_callback);
 
   // We assigned this task to a worker.
   // (Note this means that we sent the task to the worker. The assignment
@@ -2294,6 +2255,72 @@ void NodeManager::ForwardTask(
   });
 }
 
+void NodeManager::FinishAssignTask(const TaskID &task_id, Worker &worker, bool success) {
+  // Remove the ASSIGNED task from the SWAP queue.
+  Task assigned_task;
+  TaskState state;
+  if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) {
+    return;
+  }
+
+  RAY_CHECK(state == TaskState::SWAP);
+
+  if (success) {
+    auto spec = assigned_task.GetTaskSpecification();
+    // We successfully assigned the task to the worker.
+    worker.AssignTaskId(spec.TaskId());
+    worker.AssignJobId(spec.JobId());
+    // Actor tasks require extra accounting to track the actor's state.
+    if (spec.IsActorTask()) {
+      auto actor_entry = actor_registry_.find(spec.ActorId());
+      RAY_CHECK(actor_entry != actor_registry_.end());
+      // Process any new actor handles that were created since the
+      // previous task on this handle was executed. The first task
+      // submitted on a new actor handle will depend on the dummy object
+      // returned by the previous task, so the dependency will not be
+      // released until this first task is submitted.
+      for (auto &new_handle_id : spec.NewActorHandles()) {
+        // Get the execution dependency for the first task submitted on the new
+        // actor handle. Since the new actor handle was created after this task
+        // began and before this task finished, it must have the same execution
+        // dependency.
+        const auto &execution_dependencies =
+            assigned_task.GetTaskExecutionSpec().ExecutionDependencies();
+        // TODO(swang): We expect this task to have exactly 1 execution dependency,
+        // the dummy object returned by the previous actor task. However, this
+        // leaks information about the TaskExecutionSpecification implementation.
+        RAY_CHECK(execution_dependencies.size() == 1);
+        const ObjectID &execution_dependency = execution_dependencies.front();
+        // Add the new handle and give it a reference to the finished task's
+        // execution dependency.
+        actor_entry->second.AddHandle(new_handle_id, execution_dependency);
+      }
+
+      // TODO(swang): For actors with multiple actor handles, to
+      // guarantee that tasks are replayed in the same order after a
+      // failure, we must update the task's execution dependency to be
+      // the actor's current execution dependency.
+    }
+
+    // Mark the task as running.
+    // (See design_docs/task_states.rst for the state transition diagram.)
+    local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING);
+    // Notify the task dependency manager that we no longer need this task's
+    // object dependencies.
+    RAY_CHECK(task_dependency_manager_.UnsubscribeDependencies(spec.TaskId()));
+  } else {
+    RAY_LOG(WARNING) << "Failed to send task to worker, disconnecting client";
+    // We failed to send the task to the worker, so disconnect the worker.
+    ProcessDisconnectClientMessage(worker.Connection());
+    // Queue this task for future assignment. We need to do this since
+    // DispatchTasks() removed it from the ready queue. The task will be
+    // assigned to a worker once one becomes available.
+    // (See design_docs/task_states.rst for the state transition diagram.)
+    local_queues_.QueueTasks({assigned_task}, TaskState::READY);
+    DispatchTasks(MakeTasksWithResources({assigned_task}));
+  }
+}
+
 void NodeManager::DumpDebugState() const {
   std::fstream fs;
   fs.open(initial_config_.session_dir + "/debug_state.txt",

src/ray/raylet/node_manager.h

@@ -375,11 +375,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   void ProcessRegisterClientRequestMessage(
       const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data);
 
-  /// Process client message of GetTask
+  /// Handle the case that a worker is available.
   ///
-  /// \param client The client that sent the message.
+  /// \param client The connection for the worker.
   /// \return Void.
-  void ProcessGetTaskMessage(const std::shared_ptr<LocalClientConnection> &client);
+  void HandleWorkerAvailable(const std::shared_ptr<LocalClientConnection> &client);
 
   /// Handle a client that has disconnected. This can be called multiple times
   /// on the same client because this is triggered both when a client
@@ -459,6 +459,14 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   void HandleDisconnectedActor(const ActorID &actor_id, bool was_local,
                                bool intentional_disconnect);
 
+  /// Finish assigning a task to a worker.
+  ///
+  /// \param task_id Id of the task.
+  /// \param worker Worker which the task is assigned to.
+  /// \param success Whether the task is successfully assigned to the worker.
+  /// \return void.
+  void FinishAssignTask(const TaskID &task_id, Worker &worker, bool success);
+
   /// Handle a `ForwardTask` request.
   void HandleForwardTask(const rpc::ForwardTaskRequest &request,
                          rpc::ForwardTaskReply *reply,
@@ -523,7 +531,8 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   /// The RPC service.
   rpc::NodeManagerGrpcService node_manager_service_;
 
-  /// The `ClientCallManager` object that is shared by all `NodeManagerClient`s.
+  /// The `ClientCallManager` object that is shared by all `NodeManagerClient`s
+  /// as well as all `WorkerTaskClient`s.
   rpc::ClientCallManager client_call_manager_;
 
   /// Map from node ids to clients of the remote node managers.

src/ray/raylet/raylet.h

@@ -18,7 +18,6 @@ namespace raylet {
 
 using rpc::ClientTableData;
 
-class Task;
 class NodeManager;
 
 class Raylet {

src/ray/raylet/scheduling_policy.cc

@@ -51,7 +51,7 @@ std::unordered_map<TaskID, ClientID> SchedulingPolicy::Schedule(
           ResourceSet(node_resources.GetAvailableResources());
       // We have to subtract the current "load" because we set the current "load"
       // to be the resources used by tasks that are in the
-      // `SchedulingQueue::ready_queue_` in NodeManager::ProcessGetTaskMessage's
+      // `SchedulingQueue::ready_queue_` in NodeManager::HandleWorkerAvailable's
       // call to SchedulingQueue::GetResourceLoad.
       available_node_resources.SubtractResources(node_resources.GetLoadResources());
       RAY_LOG(DEBUG) << "client_id " << node_client_id

src/ray/raylet/worker.cc

@@ -11,13 +11,20 @@ namespace raylet {
 
 /// A constructor responsible for initializing the state of a worker.
 Worker::Worker(pid_t pid, const Language &language, int port,
-               std::shared_ptr<LocalClientConnection> connection)
+               std::shared_ptr<LocalClientConnection> connection,
+               rpc::ClientCallManager &client_call_manager)
     : pid_(pid),
       language_(language),
       port_(port),
       connection_(connection),
       dead_(false),
-      blocked_(false) {}
+      blocked_(false),
+      client_call_manager_(client_call_manager) {
+  if (port_ > 0) {
+    rpc_client_ = std::unique_ptr<rpc::WorkerTaskClient>(
+        new rpc::WorkerTaskClient("127.0.0.1", port_, client_call_manager_));
+  }
+}
 
 void Worker::MarkDead() { dead_ = true; }
 
@@ -103,6 +110,44 @@ void Worker::AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) {
   task_resource_ids_.Release(cpu_resources);
 }
 
+bool Worker::UsePush() const { return rpc_client_ != nullptr; }
+
+void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
+                        const std::function<void(Status)> finish_assign_callback) {
+  const TaskSpecification &spec = task.GetTaskSpecification();
+  if (rpc_client_ != nullptr) {
+    // Use push mode.
+    RAY_CHECK(port_ > 0);
+    rpc::AssignTaskRequest request;
+    request.mutable_task()->mutable_task_spec()->CopyFrom(
+        task.GetTaskSpecification().GetMessage());
+    request.mutable_task()->mutable_task_execution_spec()->CopyFrom(
+        task.GetTaskExecutionSpec().GetMessage());
+    request.set_resource_ids(resource_id_set.Serialize());
+
+    auto status = rpc_client_->AssignTask(
+        request, [](Status status, const rpc::AssignTaskReply &reply) {
+          // Worker has finished this task. There's nothing to do here
+          // and assigning new task will be done when raylet receives
+          // `TaskDone` message.
+        });
+    finish_assign_callback(status);
+  } else {
+    // Use pull mode. This corresponds to existing python/java workers that haven't been
+    // migrated to core worker architecture.
+    flatbuffers::FlatBufferBuilder fbb;
+    auto resource_id_set_flatbuf = resource_id_set.ToFlatbuf(fbb);
+
+    auto message =
+        protocol::CreateGetTaskReply(fbb, fbb.CreateString(spec.Serialize()),
+                                     fbb.CreateVector(resource_id_set_flatbuf));
+    fbb.Finish(message);
+    Connection()->WriteMessageAsync(
+        static_cast<int64_t>(protocol::MessageType::ExecuteTask), fbb.GetSize(),
+        fbb.GetBufferPointer(), finish_assign_callback);
+  }
+}
+
 }  // namespace raylet
 
 }  // end namespace ray

src/ray/raylet/worker.h

@@ -6,7 +6,9 @@
 #include "ray/common/client_connection.h"
 #include "ray/common/id.h"
 #include "ray/common/task/scheduling_resources.h"
+#include "ray/common/task/task.h"
 #include "ray/common/task/task_common.h"
+#include "ray/rpc/worker/worker_client.h"
 
 namespace ray {
 
@@ -19,7 +21,8 @@ class Worker {
  public:
   /// A constructor that initializes a worker object.
   Worker(pid_t pid, const Language &language, int port,
-         std::shared_ptr<LocalClientConnection> connection);
+         std::shared_ptr<LocalClientConnection> connection,
+         rpc::ClientCallManager &client_call_manager);
   /// A destructor responsible for freeing all worker state.
   ~Worker() {}
   void MarkDead();
@@ -53,6 +56,10 @@ class Worker {
   ResourceIdSet ReleaseTaskCpuResources();
   void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources);
 
+  bool UsePush() const;
+  void AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
+                  const std::function<void(Status)> finish_assign_callback);
+
  private:
   /// The worker's PID.
   pid_t pid_;
@@ -81,6 +88,11 @@ class Worker {
   // of a task.
   ResourceIdSet task_resource_ids_;
   std::unordered_set<TaskID> blocked_task_ids_;
+  /// The `ClientCallManager` object that is shared by `WorkerTaskClient` from all
+  /// workers.
+  rpc::ClientCallManager &client_call_manager_;
+  /// The rpc client to send tasks to this worker.
+  std::unique_ptr<rpc::WorkerTaskClient> rpc_client_;
 };
 
 }  // namespace raylet

src/ray/raylet/worker_pool_test.cc

@@ -70,7 +70,11 @@ class WorkerPoolMock : public WorkerPool {
 
 class WorkerPoolTest : public ::testing::Test {
  public:
-  WorkerPoolTest() : worker_pool_(), io_service_(), error_message_type_(1) {}
+  WorkerPoolTest()
+      : worker_pool_(),
+        io_service_(),
+        error_message_type_(1),
+        client_call_manager_(io_service_) {}
 
   std::shared_ptr<Worker> CreateWorker(pid_t pid,
                                        const Language &language = Language::PYTHON) {
@@ -85,7 +89,8 @@ class WorkerPoolTest : public ::testing::Test {
     auto client =
         LocalClientConnection::Create(client_handler, message_handler, std::move(socket),
                                       "worker", {}, error_message_type_);
-    return std::shared_ptr<Worker>(new Worker(pid, language, -1, client));
+    return std::shared_ptr<Worker>(
+        new Worker(pid, language, -1, client, client_call_manager_));
   }
 
   void SetWorkerCommands(const WorkerCommandMap &worker_commands) {
@@ -97,6 +102,7 @@ class WorkerPoolTest : public ::testing::Test {
   WorkerPoolMock worker_pool_;
   boost::asio::io_service io_service_;
   int64_t error_message_type_;
+  rpc::ClientCallManager client_call_manager_;
 
  private:
   void HandleNewClient(LocalClientConnection &){};