Replace ReturnIds with NumReturns in TaskInfo to reduce the size (#4854)

* Refine TaskInfo

* Fix

* Add a test to print task info size

* Lint

* Refine

src/ray/gcs/format/gcs.fbs

@@ -89,9 +89,8 @@ table TaskInfo {
   new_actor_handles: string;
   // Task arguments.
   args: [Arg];
-  // Object IDs of return values. This is a long string that concatenate
-  // all of the return object IDs of this task.
-  returns: string;
+  // Number of return objects.
+  num_returns: int;
   // The required_resources vector indicates the quantities of the different
   // resources required by this task.
   required_resources: [ResourcePair];

src/ray/raylet/task_spec.cc

@@ -92,20 +92,14 @@ TaskSpecification::TaskSpecification(
     arguments.push_back(argument->ToFlatbuffer(fbb));
   }
 
-  // Generate return ids.
-  std::vector<ray::ObjectID> returns;
-  for (int64_t i = 1; i < num_returns + 1; ++i) {
-    returns.push_back(ObjectID::for_task_return(task_id, i));
-  }
-
   // Serialize the TaskSpecification.
   auto spec = CreateTaskInfo(
       fbb, to_flatbuf(fbb, driver_id), to_flatbuf(fbb, task_id),
       to_flatbuf(fbb, parent_task_id), parent_counter, to_flatbuf(fbb, actor_creation_id),
       to_flatbuf(fbb, actor_creation_dummy_object_id), max_actor_reconstructions,
       to_flatbuf(fbb, actor_id), to_flatbuf(fbb, actor_handle_id), actor_counter,
-      ids_to_flatbuf(fbb, new_actor_handles), fbb.CreateVector(arguments),
-      ids_to_flatbuf(fbb, returns), map_to_flatbuf(fbb, required_resources),
+      ids_to_flatbuf(fbb, new_actor_handles), fbb.CreateVector(arguments), num_returns,
+      map_to_flatbuf(fbb, required_resources),
       map_to_flatbuf(fbb, required_placement_resources), language,
       string_vec_to_flatbuf(fbb, function_descriptor));
   fbb.Finish(spec);
@@ -167,12 +161,12 @@ int64_t TaskSpecification::NumArgs() const {
 
 int64_t TaskSpecification::NumReturns() const {
   auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
-  return (message->returns()->size() / kUniqueIDSize);
+  return message->num_returns();
 }
 
 ObjectID TaskSpecification::ReturnId(int64_t return_index) const {
   auto message = flatbuffers::GetRoot<TaskInfo>(spec_.data());
-  return ids_from_flatbuf<ObjectID>(*message->returns())[return_index];
+  return ObjectID::for_task_return(TaskId(), return_index + 1);
 }
 
 bool TaskSpecification::ArgByRef(int64_t arg_index) const {

src/ray/raylet/task_test.cc

@@ -1,5 +1,6 @@
 #include "gtest/gtest.h"
 
+#include "ray/common/common_protocol.h"
 #include "ray/raylet/task_spec.h"
 
 namespace ray {
@@ -47,6 +48,56 @@ TEST(IdPropertyTest, TestIdProperty) {
   ASSERT_TRUE(ObjectID::nil().is_nil());
 }
 
+TEST(TaskSpecTest, TaskInfoSize) {
+  std::vector<ObjectID> references = {ObjectID::from_random(), ObjectID::from_random()};
+  auto arguments_1 = std::make_shared<TaskArgumentByReference>(references);
+  std::string one_arg("This is an value argument.");
+  auto arguments_2 = std::make_shared<TaskArgumentByValue>(
+      reinterpret_cast<const uint8_t *>(one_arg.c_str()), one_arg.size());
+  std::vector<std::shared_ptr<TaskArgument>> task_arguments({arguments_1, arguments_2});
+  auto task_id = TaskID::from_random();
+  {
+    flatbuffers::FlatBufferBuilder fbb;
+    std::vector<flatbuffers::Offset<Arg>> arguments;
+    for (auto &argument : task_arguments) {
+      arguments.push_back(argument->ToFlatbuffer(fbb));
+    }
+    // General task.
+    auto spec = CreateTaskInfo(
+        fbb, to_flatbuf(fbb, DriverID::from_random()), to_flatbuf(fbb, task_id),
+        to_flatbuf(fbb, TaskID::from_random()), 0, to_flatbuf(fbb, ActorID::nil()),
+        to_flatbuf(fbb, ObjectID::nil()), 0, to_flatbuf(fbb, ActorID::nil()),
+        to_flatbuf(fbb, ActorHandleID::nil()), 0,
+        ids_to_flatbuf(fbb, std::vector<ObjectID>()), fbb.CreateVector(arguments), 1,
+        map_to_flatbuf(fbb, {}), map_to_flatbuf(fbb, {}), Language::PYTHON,
+        string_vec_to_flatbuf(fbb, {"PackageName", "ClassName", "FunctionName"}));
+    fbb.Finish(spec);
+    RAY_LOG(ERROR) << "Ordinary task info size: " << fbb.GetSize();
+  }
+
+  {
+    flatbuffers::FlatBufferBuilder fbb;
+    std::vector<flatbuffers::Offset<Arg>> arguments;
+    for (auto &argument : task_arguments) {
+      arguments.push_back(argument->ToFlatbuffer(fbb));
+    }
+    // General task.
+    auto spec = CreateTaskInfo(
+        fbb, to_flatbuf(fbb, DriverID::from_random()), to_flatbuf(fbb, task_id),
+        to_flatbuf(fbb, TaskID::from_random()), 10,
+        to_flatbuf(fbb, ActorID::from_random()), to_flatbuf(fbb, ObjectID::from_random()),
+        10000000, to_flatbuf(fbb, ActorID::from_random()),
+        to_flatbuf(fbb, ActorHandleID::from_random()), 20,
+        ids_to_flatbuf(fbb, std::vector<ObjectID>(
+                                {ObjectID::from_random(), ObjectID::from_random()})),
+        fbb.CreateVector(arguments), 2, map_to_flatbuf(fbb, {}), map_to_flatbuf(fbb, {}),
+        Language::PYTHON,
+        string_vec_to_flatbuf(fbb, {"PackageName", "ClassName", "FunctionName"}));
+    fbb.Finish(spec);
+    RAY_LOG(ERROR) << "Actor task info size: " << fbb.GetSize();
+  }
+}
+
 }  // namespace raylet
 
 }  // namespace ray