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

* Refine TaskInfo

* Fix

* Add a test to print task info size

* Lint

* Refine

java/runtime/src/main/java/org/ray/runtime/AbstractRayRuntime.java

@@ -390,7 +390,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
         actor.increaseTaskCounter(),
         actor.getNewActorHandles().toArray(new UniqueId[0]),
         ArgumentsBuilder.wrap(args, language == TaskLanguage.PYTHON),
-        returnIds,
+        numReturns,
         resources,
         language,
         functionDescriptor

java/runtime/src/main/java/org/ray/runtime/raylet/RayletClientImpl.java

@@ -154,6 +154,7 @@ public class RayletClientImpl implements RayletClient {
     UniqueId actorId = UniqueId.fromByteBuffer(info.actorIdAsByteBuffer());
     UniqueId actorHandleId = UniqueId.fromByteBuffer(info.actorHandleIdAsByteBuffer());
     int actorCounter = info.actorCounter();
+    int numReturns = info.numReturns();
 
     // Deserialize new actor handles
     UniqueId[] newActorHandles = IdUtil.getUniqueIdsFromByteBuffer(
@@ -177,8 +178,6 @@ public class RayletClientImpl implements RayletClient {
         args[i] = FunctionArg.passByValue(data);
       }
     }
-    // Deserialize return ids
-    ObjectId[] returnIds = IdUtil.getObjectIdsFromByteBuffer(info.returnsAsByteBuffer());
 
     // Deserialize required resources;
     Map<String, Double> resources = new HashMap<>();
@@ -193,7 +192,7 @@ public class RayletClientImpl implements RayletClient {
     );
     return new TaskSpec(driverId, taskId, parentTaskId, parentCounter, actorCreationId,
         maxActorReconstructions, actorId, actorHandleId, actorCounter, newActorHandles,
-        args, returnIds, resources, TaskLanguage.JAVA, functionDescriptor);
+        args, numReturns, resources, TaskLanguage.JAVA, functionDescriptor);
   }
 
   private static ByteBuffer convertTaskSpecToFlatbuffer(TaskSpec task) {
@@ -211,6 +210,7 @@ public class RayletClientImpl implements RayletClient {
     final int actorIdOffset = fbb.createString(task.actorId.toByteBuffer());
     final int actorHandleIdOffset = fbb.createString(task.actorHandleId.toByteBuffer());
     final int actorCounter = task.actorCounter;
+    final int numReturnsOffset = task.numReturns;
 
     // Serialize the new actor handles.
     int newActorHandlesOffset
@@ -234,9 +234,6 @@ public class RayletClientImpl implements RayletClient {
     }
     int argsOffset = fbb.createVectorOfTables(argsOffsets);
 
-    // Serialize returns
-    int returnsOffset = fbb.createString(IdUtil.concatIds(task.returnIds));
-
     // Serialize required resources
     // The required_resources vector indicates the quantities of the different
     // resources required by this task. The index in this vector corresponds to
@@ -292,7 +289,7 @@ public class RayletClientImpl implements RayletClient {
         actorCounter,
         newActorHandlesOffset,
         argsOffset,
-        returnsOffset,
+        numReturnsOffset,
         requiredResourcesOffset,
         requiredPlacementResourcesOffset,
         language,

java/runtime/src/main/java/org/ray/runtime/task/TaskSpec.java

@@ -11,6 +11,7 @@ import org.ray.api.id.UniqueId;
 import org.ray.runtime.functionmanager.FunctionDescriptor;
 import org.ray.runtime.functionmanager.JavaFunctionDescriptor;
 import org.ray.runtime.functionmanager.PyFunctionDescriptor;
+import org.ray.runtime.util.IdUtil;
 
 /**
  * Represents necessary information of a task for scheduling and executing.
@@ -50,7 +51,10 @@ public class TaskSpec {
   // Task arguments.
   public final FunctionArg[] args;
 
-  // return ids
+  // number of return objects.
+  public final int numReturns;
+
+  // returns ids.
   public final ObjectId[] returnIds;
 
   // The task's resource demands.
@@ -86,7 +90,7 @@ public class TaskSpec {
       int actorCounter,
       UniqueId[] newActorHandles,
       FunctionArg[] args,
-      ObjectId[] returnIds,
+      int numReturns,
       Map<String, Double> resources,
       TaskLanguage language,
       FunctionDescriptor functionDescriptor) {
@@ -101,7 +105,11 @@ public class TaskSpec {
     this.actorCounter = actorCounter;
     this.newActorHandles = newActorHandles;
     this.args = args;
-    this.returnIds = returnIds;
+    this.numReturns = numReturns;
+    returnIds = new ObjectId[numReturns];
+    for (int i = 0; i < numReturns; ++i) {
+      returnIds[i] = IdUtil.computeReturnId(taskId, i + 1);
+    }
     this.resources = resources;
     this.language = language;
     if (language == TaskLanguage.JAVA) {
@@ -145,7 +153,7 @@ public class TaskSpec {
         ", actorCounter=" + actorCounter +
         ", newActorHandles=" + Arrays.toString(newActorHandles) +
         ", args=" + Arrays.toString(args) +
-        ", returnIds=" + Arrays.toString(returnIds) +
+        ", numReturns=" + numReturns +
         ", resources=" + resources +
         ", language=" + language +
         ", functionDescriptor=" + functionDescriptor +