[Core worker] Task execution passes TaskInfo struct to executor (#5032)

src/ray/core_worker/common.h

@@ -68,6 +68,18 @@ class TaskArg {
   const std::shared_ptr<Buffer> data_;
 };
 
+enum class TaskType { NORMAL_TASK, ACTOR_CREATION_TASK, ACTOR_TASK };
+
+/// Information of a task
+struct TaskInfo {
+  /// The ID of task.
+  const TaskID task_id;
+  /// The driver ID.
+  const DriverID driver_id;
+  /// The type of task.
+  const TaskType task_type;
+};
+
 /// Task specification, which includes the immutable information about the task
 /// which are determined at the submission time.
 /// TODO(zhijunfu): this can be removed after everything is moved to protobuf.

src/ray/core_worker/mock_worker.cc

@@ -22,7 +22,7 @@ class MockWorker {
   void Run() {
     auto executor_func = [this](const RayFunction &ray_function,
                                 const std::vector<std::shared_ptr<Buffer>> &args,
-                                const TaskID &task_id, int num_returns) {
+                                const TaskInfo &task_info, int num_returns) {
       // Note that this doesn't include dummy object id.
       RAY_CHECK(num_returns >= 0);
 
@@ -36,7 +36,7 @@ class MockWorker {
 
       // Write the merged content to each of return ids.
       for (int i = 0; i < num_returns; i++) {
-        ObjectID id = ObjectID::ForTaskReturn(task_id, i + 1);
+        ObjectID id = ObjectID::ForTaskReturn(task_info.task_id, i + 1);
         RAY_CHECK_OK(worker_.Objects().Put(memory_buffer, id));
       }
       return Status::OK();

src/ray/core_worker/task_execution.cc

@@ -37,6 +37,17 @@ Status CoreWorkerTaskExecutionInterface::Run(const TaskExecutor &executor) {
       std::vector<std::shared_ptr<Buffer>> 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.DriverId(), task_type};
+
       auto num_returns = spec.NumReturns();
       if (spec.IsActorCreationTask() || spec.IsActorTask()) {
         RAY_CHECK(num_returns > 0);
@@ -44,7 +55,7 @@ Status CoreWorkerTaskExecutionInterface::Run(const TaskExecutor &executor) {
         num_returns--;
       }
 
-      status = executor(func, args, spec.TaskId(), num_returns);
+      status = executor(func, args, task_info, num_returns);
       // TODO(zhijunfu):
       // 1. Check and handle failure.
       // 2. Save or load checkpoint.

src/ray/core_worker/task_execution.h

@@ -26,7 +26,7 @@ class CoreWorkerTaskExecutionInterface {
   /// \return Status.
   using TaskExecutor = std::function<Status(
       const RayFunction &ray_function, const std::vector<std::shared_ptr<Buffer>> &args,
-      const TaskID &task_id, int num_returns)>;
+      const TaskInfo &task_info, int num_returns)>;
 
   /// Start receving and executes tasks in a infinite loop.
   /// \return Status.