[direct call] changes raylet to push tasks to worker

src/ray/core_worker/core_worker.cc

@@ -5,7 +5,8 @@ namespace ray {
 
 CoreWorker::CoreWorker(const enum WorkerType worker_type, const ::Language language,
                        const std::string &store_socket, const std::string &raylet_socket,
-                       const JobID &job_id)
+                       const JobID &job_id,
+                       CoreWorkerTaskExecutionInterface::TaskExecutor execution_callback)
     : worker_type_(worker_type),
       language_(language),
       raylet_socket_(raylet_socket),
@@ -14,9 +15,10 @@ CoreWorker::CoreWorker(const enum WorkerType worker_type, const ::Language langu
       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

@@ -25,7 +25,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,
+             CoreWorkerTaskExecutionInterface::TaskExecutor execution_callback = nullptr);
 
   /// Type of this worker.
   enum WorkerType WorkerType() const { return worker_type_; }

src/ray/core_worker/mock_worker.cc

@@ -16,11 +16,7 @@ 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)
-      : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
-                JobID::FromRandom()) {}
-
-  void Run() {
+  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) {
@@ -40,17 +36,23 @@ class MockWorker {
       // 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));
+        RAY_CHECK_OK(worker_->Objects().Put(return_value, id));
       }
       return Status::OK();
     };
 
+    worker_ = std::unique_ptr<CoreWorker>(
+        new CoreWorker(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
+                       JobID::FromRandom(), executor_func));
+  }
+
+  void Run() {
     // Start executing tasks.
-    worker_.Execution().Run(executor_func);
+    worker_->Execution().Run();
   }
 
  private:
-  CoreWorker worker_;
+  std::unique_ptr<CoreWorker> worker_;
 };
 
 }  // namespace ray

src/ray/core_worker/task_execution.cc

@@ -7,65 +7,63 @@ 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_)));
+      std::unique_ptr<CoreWorkerRayletTaskReceiver>(
+          new CoreWorkerRayletTaskReceiver(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 raylet::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()};
-
-      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.
-    }
+  RayFunction func{spec.GetLanguage(), spec.FunctionDescriptor()};
+
+  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--;
   }
 
+  auto status = execution_callback_(func, args, task_info, num_returns);
+  // TODO(zhijunfu):
+  // 1. Check and handle failure.
+  // 2. Save or load checkpoint.
+  return status;
+}
+
+Status CoreWorkerTaskExecutionInterface::Run() {
+  // Run main IO service.
+  main_service_.run();
+
   // should never reach here.
   return Status::OK();
 }

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);
+  Status Run();
 
  private:
   /// Build arguments for task executor. This would loop through all the arguments
@@ -53,11 +54,17 @@ class CoreWorkerTaskExecutionInterface {
   Status BuildArgsForExecutor(const raylet::TaskSpecification &spec,
                               std::vector<std::shared_ptr<RayObject>> *args);
 
+  /// Execute a task.
+  Status ExecuteTask(const raylet::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,14 @@ 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<raylet::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) {
+    boost::asio::io_service &io_service, rpc::GrpcServer &server,
+    const TaskHandler &task_handler)
+    : task_service_(io_service, *this), task_handler_(task_handler) {
   server.RegisterService(task_service_);
 }
 

src/ray/core_worker/transport/raylet_transport.h

@@ -31,17 +31,9 @@ class CoreWorkerRayletTaskSubmitter : public CoreWorkerTaskSubmitter {
 class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver,
                                      public rpc::WorkerTaskHandler {
  public:
-  CoreWorkerRayletTaskReceiver(std::unique_ptr<RayletClient> &raylet_client,
-                               boost::asio::io_service &io_service,
-                               rpc::GrpcServer &server);
+  CoreWorkerRayletTaskReceiver(boost::asio::io_service &io_service,
+                               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
   /// `done_callback` should be called.
@@ -54,12 +46,10 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver,
                         rpc::RequestDoneCallback done_callback) override;
 
  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 raylet::TaskSpecification &task_spec)>;
-
-  // Get tasks for execution.
-  virtual Status GetTasks(std::vector<TaskSpec> *tasks) = 0;
 };
 
 }  // namespace ray

src/ray/protobuf/worker.proto

@@ -5,12 +5,16 @@ package ray.rpc;
 message AssignTaskRequest {
   // The ID of the task to be pushed.
   bytes task_id = 1;
-  // The task to be pushed. This should include task_id.
+  // 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;
+  // 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;
 }
 
 message AssignTaskReply {