From a39982e676658b71305a47c1f026f3a0efb388b2 Mon Sep 17 00:00:00 2001
From: Kai Yang <kfstorm@outlook.com>
Date: Fri, 28 Jun 2019 10:59:45 +0800
Subject: [PATCH] [Core worker] Task execution passes TaskInfo struct to
 executor (#5032)

---
 src/ray/core_worker/common.h          | 12 ++++++++++++
 src/ray/core_worker/mock_worker.cc    |  4 ++--
 src/ray/core_worker/task_execution.cc | 13 ++++++++++++-
 src/ray/core_worker/task_execution.h  |  2 +-
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/ray/core_worker/common.h b/src/ray/core_worker/common.h
index 3fda40661..b5288c033 100644
--- a/src/ray/core_worker/common.h
+++ b/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.
diff --git a/src/ray/core_worker/mock_worker.cc b/src/ray/core_worker/mock_worker.cc
index b7f4b5a63..eececbe4f 100644
--- a/src/ray/core_worker/mock_worker.cc
+++ b/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();
diff --git a/src/ray/core_worker/task_execution.cc b/src/ray/core_worker/task_execution.cc
index 701ae3124..68ab69299 100644
--- a/src/ray/core_worker/task_execution.cc
+++ b/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.
diff --git a/src/ray/core_worker/task_execution.h b/src/ray/core_worker/task_execution.h
index f4b44b9e1..f3e78f2f9 100644
--- a/src/ray/core_worker/task_execution.h
+++ b/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.