[xray] Implement Actor Reconstruction (#3332)

* Implement Actor Reconstruction

* fix

* fix actor handle __del__

* fix lint

* add comment

* Remove actorCreationDummyObjectId

* address comments

* fix

* address comments

* avoid copy

* change log to debug

* fix error name

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

@@ -270,6 +270,10 @@ public abstract class AbstractRayRuntime implements RayRuntime {
       resources.put(ResourceUtil.CPU_LITERAL, 0.0);
     }
 
+    int maxActorReconstruction = 0;
+    if (taskOptions instanceof ActorCreationOptions) {
+      maxActorReconstruction = ((ActorCreationOptions) taskOptions).maxReconstructions;
+    }
     RayFunction rayFunction = functionManager.getFunction(current.driverId, func);
     return new TaskSpec(
         current.driverId,
@@ -277,6 +281,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
         current.taskId,
         -1,
         actorCreationId,
+        maxActorReconstruction,
         actor.getId(),
         actor.getHandleId(),
         actor.increaseTaskCounter(),

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

@@ -67,7 +67,6 @@ public final class RayActorImpl<T> implements RayActor<T>, Externalizable {
     return taskCounter++;
   }
 
-
   private UniqueId computeNextActorHandleId() {
     byte[] bytes = Sha1Digestor.digest(handleId.getBytes(), ++numForks);
     return new UniqueId(bytes);

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

@@ -79,6 +79,7 @@ public class WorkerContext {
         UniqueId.NIL,
         0,
         UniqueId.NIL,
+        0,
         UniqueId.NIL,
         UniqueId.NIL,
         0,

java/runtime/src/main/java/org/ray/runtime/generated/TaskInfo.java

@@ -29,31 +29,32 @@ public final class TaskInfo extends Table {
   public String actorCreationDummyObjectId() { int o = __offset(14); return o != 0 ? __string(o + bb_pos) : null; }
   public ByteBuffer actorCreationDummyObjectIdAsByteBuffer() { return __vector_as_bytebuffer(14, 1); }
   public ByteBuffer actorCreationDummyObjectIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 14, 1); }
-  public String actorId() { int o = __offset(16); return o != 0 ? __string(o + bb_pos) : null; }
-  public ByteBuffer actorIdAsByteBuffer() { return __vector_as_bytebuffer(16, 1); }
-  public ByteBuffer actorIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 16, 1); }
-  public String actorHandleId() { int o = __offset(18); return o != 0 ? __string(o + bb_pos) : null; }
-  public ByteBuffer actorHandleIdAsByteBuffer() { return __vector_as_bytebuffer(18, 1); }
-  public ByteBuffer actorHandleIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 18, 1); }
-  public int actorCounter() { int o = __offset(20); return o != 0 ? bb.getInt(o + bb_pos) : 0; }
-  public boolean isActorCheckpointMethod() { int o = __offset(22); return o != 0 ? 0!=bb.get(o + bb_pos) : false; }
-  public String functionId() { int o = __offset(24); return o != 0 ? __string(o + bb_pos) : null; }
-  public ByteBuffer functionIdAsByteBuffer() { return __vector_as_bytebuffer(24, 1); }
-  public ByteBuffer functionIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 24, 1); }
+  public int maxActorReconstructions() { int o = __offset(16); return o != 0 ? bb.getInt(o + bb_pos) : 0; }
+  public String actorId() { int o = __offset(18); return o != 0 ? __string(o + bb_pos) : null; }
+  public ByteBuffer actorIdAsByteBuffer() { return __vector_as_bytebuffer(18, 1); }
+  public ByteBuffer actorIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 18, 1); }
+  public String actorHandleId() { int o = __offset(20); return o != 0 ? __string(o + bb_pos) : null; }
+  public ByteBuffer actorHandleIdAsByteBuffer() { return __vector_as_bytebuffer(20, 1); }
+  public ByteBuffer actorHandleIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 20, 1); }
+  public int actorCounter() { int o = __offset(22); return o != 0 ? bb.getInt(o + bb_pos) : 0; }
+  public boolean isActorCheckpointMethod() { int o = __offset(24); return o != 0 ? 0!=bb.get(o + bb_pos) : false; }
+  public String functionId() { int o = __offset(26); return o != 0 ? __string(o + bb_pos) : null; }
+  public ByteBuffer functionIdAsByteBuffer() { return __vector_as_bytebuffer(26, 1); }
+  public ByteBuffer functionIdInByteBuffer(ByteBuffer _bb) { return __vector_in_bytebuffer(_bb, 26, 1); }
   public Arg args(int j) { return args(new Arg(), j); }
-  public Arg args(Arg obj, int j) { int o = __offset(26); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; }
-  public int argsLength() { int o = __offset(26); return o != 0 ? __vector_len(o) : 0; }
-  public String returns(int j) { int o = __offset(28); return o != 0 ? __string(__vector(o) + j * 4) : null; }
-  public int returnsLength() { int o = __offset(28); return o != 0 ? __vector_len(o) : 0; }
+  public Arg args(Arg obj, int j) { int o = __offset(28); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; }
+  public int argsLength() { int o = __offset(28); return o != 0 ? __vector_len(o) : 0; }
+  public String returns(int j) { int o = __offset(30); return o != 0 ? __string(__vector(o) + j * 4) : null; }
+  public int returnsLength() { int o = __offset(30); return o != 0 ? __vector_len(o) : 0; }
   public ResourcePair requiredResources(int j) { return requiredResources(new ResourcePair(), j); }
-  public ResourcePair requiredResources(ResourcePair obj, int j) { int o = __offset(30); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; }
-  public int requiredResourcesLength() { int o = __offset(30); return o != 0 ? __vector_len(o) : 0; }
+  public ResourcePair requiredResources(ResourcePair obj, int j) { int o = __offset(32); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; }
+  public int requiredResourcesLength() { int o = __offset(32); return o != 0 ? __vector_len(o) : 0; }
   public ResourcePair requiredPlacementResources(int j) { return requiredPlacementResources(new ResourcePair(), j); }
-  public ResourcePair requiredPlacementResources(ResourcePair obj, int j) { int o = __offset(32); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; }
-  public int requiredPlacementResourcesLength() { int o = __offset(32); return o != 0 ? __vector_len(o) : 0; }
-  public int language() { int o = __offset(34); return o != 0 ? bb.getInt(o + bb_pos) : 0; }
-  public String functionDescriptor(int j) { int o = __offset(36); return o != 0 ? __string(__vector(o) + j * 4) : null; }
-  public int functionDescriptorLength() { int o = __offset(36); return o != 0 ? __vector_len(o) : 0; }
+  public ResourcePair requiredPlacementResources(ResourcePair obj, int j) { int o = __offset(34); return o != 0 ? obj.__assign(__indirect(__vector(o) + j * 4), bb) : null; }
+  public int requiredPlacementResourcesLength() { int o = __offset(34); return o != 0 ? __vector_len(o) : 0; }
+  public int language() { int o = __offset(36); return o != 0 ? bb.getInt(o + bb_pos) : 0; }
+  public String functionDescriptor(int j) { int o = __offset(38); return o != 0 ? __string(__vector(o) + j * 4) : null; }
+  public int functionDescriptorLength() { int o = __offset(38); return o != 0 ? __vector_len(o) : 0; }
 
   public static int createTaskInfo(FlatBufferBuilder builder,
       int driver_idOffset,
@@ -62,6 +63,7 @@ public final class TaskInfo extends Table {
       int parent_counter,
       int actor_creation_idOffset,
       int actor_creation_dummy_object_idOffset,
+      int max_actor_reconstructions,
       int actor_idOffset,
       int actor_handle_idOffset,
       int actor_counter,
@@ -73,7 +75,7 @@ public final class TaskInfo extends Table {
       int required_placement_resourcesOffset,
       int language,
       int function_descriptorOffset) {
-    builder.startObject(17);
+    builder.startObject(18);
     TaskInfo.addFunctionDescriptor(builder, function_descriptorOffset);
     TaskInfo.addLanguage(builder, language);
     TaskInfo.addRequiredPlacementResources(builder, required_placement_resourcesOffset);
@@ -84,6 +86,7 @@ public final class TaskInfo extends Table {
     TaskInfo.addActorCounter(builder, actor_counter);
     TaskInfo.addActorHandleId(builder, actor_handle_idOffset);
     TaskInfo.addActorId(builder, actor_idOffset);
+    TaskInfo.addMaxActorReconstructions(builder, max_actor_reconstructions);
     TaskInfo.addActorCreationDummyObjectId(builder, actor_creation_dummy_object_idOffset);
     TaskInfo.addActorCreationId(builder, actor_creation_idOffset);
     TaskInfo.addParentCounter(builder, parent_counter);
@@ -94,32 +97,33 @@ public final class TaskInfo extends Table {
     return TaskInfo.endTaskInfo(builder);
   }
 
-  public static void startTaskInfo(FlatBufferBuilder builder) { builder.startObject(17); }
+  public static void startTaskInfo(FlatBufferBuilder builder) { builder.startObject(18); }
   public static void addDriverId(FlatBufferBuilder builder, int driverIdOffset) { builder.addOffset(0, driverIdOffset, 0); }
   public static void addTaskId(FlatBufferBuilder builder, int taskIdOffset) { builder.addOffset(1, taskIdOffset, 0); }
   public static void addParentTaskId(FlatBufferBuilder builder, int parentTaskIdOffset) { builder.addOffset(2, parentTaskIdOffset, 0); }
   public static void addParentCounter(FlatBufferBuilder builder, int parentCounter) { builder.addInt(3, parentCounter, 0); }
   public static void addActorCreationId(FlatBufferBuilder builder, int actorCreationIdOffset) { builder.addOffset(4, actorCreationIdOffset, 0); }
   public static void addActorCreationDummyObjectId(FlatBufferBuilder builder, int actorCreationDummyObjectIdOffset) { builder.addOffset(5, actorCreationDummyObjectIdOffset, 0); }
-  public static void addActorId(FlatBufferBuilder builder, int actorIdOffset) { builder.addOffset(6, actorIdOffset, 0); }
-  public static void addActorHandleId(FlatBufferBuilder builder, int actorHandleIdOffset) { builder.addOffset(7, actorHandleIdOffset, 0); }
-  public static void addActorCounter(FlatBufferBuilder builder, int actorCounter) { builder.addInt(8, actorCounter, 0); }
-  public static void addIsActorCheckpointMethod(FlatBufferBuilder builder, boolean isActorCheckpointMethod) { builder.addBoolean(9, isActorCheckpointMethod, false); }
-  public static void addFunctionId(FlatBufferBuilder builder, int functionIdOffset) { builder.addOffset(10, functionIdOffset, 0); }
-  public static void addArgs(FlatBufferBuilder builder, int argsOffset) { builder.addOffset(11, argsOffset, 0); }
+  public static void addMaxActorReconstructions(FlatBufferBuilder builder, int maxActorReconstructions) { builder.addInt(6, maxActorReconstructions, 0); }
+  public static void addActorId(FlatBufferBuilder builder, int actorIdOffset) { builder.addOffset(7, actorIdOffset, 0); }
+  public static void addActorHandleId(FlatBufferBuilder builder, int actorHandleIdOffset) { builder.addOffset(8, actorHandleIdOffset, 0); }
+  public static void addActorCounter(FlatBufferBuilder builder, int actorCounter) { builder.addInt(9, actorCounter, 0); }
+  public static void addIsActorCheckpointMethod(FlatBufferBuilder builder, boolean isActorCheckpointMethod) { builder.addBoolean(10, isActorCheckpointMethod, false); }
+  public static void addFunctionId(FlatBufferBuilder builder, int functionIdOffset) { builder.addOffset(11, functionIdOffset, 0); }
+  public static void addArgs(FlatBufferBuilder builder, int argsOffset) { builder.addOffset(12, argsOffset, 0); }
   public static int createArgsVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); }
   public static void startArgsVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); }
-  public static void addReturns(FlatBufferBuilder builder, int returnsOffset) { builder.addOffset(12, returnsOffset, 0); }
+  public static void addReturns(FlatBufferBuilder builder, int returnsOffset) { builder.addOffset(13, returnsOffset, 0); }
   public static int createReturnsVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); }
   public static void startReturnsVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); }
-  public static void addRequiredResources(FlatBufferBuilder builder, int requiredResourcesOffset) { builder.addOffset(13, requiredResourcesOffset, 0); }
+  public static void addRequiredResources(FlatBufferBuilder builder, int requiredResourcesOffset) { builder.addOffset(14, requiredResourcesOffset, 0); }
   public static int createRequiredResourcesVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); }
   public static void startRequiredResourcesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); }
-  public static void addRequiredPlacementResources(FlatBufferBuilder builder, int requiredPlacementResourcesOffset) { builder.addOffset(14, requiredPlacementResourcesOffset, 0); }
+  public static void addRequiredPlacementResources(FlatBufferBuilder builder, int requiredPlacementResourcesOffset) { builder.addOffset(15, requiredPlacementResourcesOffset, 0); }
   public static int createRequiredPlacementResourcesVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); }
   public static void startRequiredPlacementResourcesVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); }
-  public static void addLanguage(FlatBufferBuilder builder, int language) { builder.addInt(15, language, 0); }
-  public static void addFunctionDescriptor(FlatBufferBuilder builder, int functionDescriptorOffset) { builder.addOffset(16, functionDescriptorOffset, 0); }
+  public static void addLanguage(FlatBufferBuilder builder, int language) { builder.addInt(16, language, 0); }
+  public static void addFunctionDescriptor(FlatBufferBuilder builder, int functionDescriptorOffset) { builder.addOffset(17, functionDescriptorOffset, 0); }
   public static int createFunctionDescriptorVector(FlatBufferBuilder builder, int[] data) { builder.startVector(4, data.length, 4); for (int i = data.length - 1; i >= 0; i--) builder.addOffset(data[i]); return builder.endVector(); }
   public static void startFunctionDescriptorVector(FlatBufferBuilder builder, int numElems) { builder.startVector(4, numElems, 4); }
   public static int endTaskInfo(FlatBufferBuilder builder) {
@@ -130,7 +134,7 @@ public final class TaskInfo extends Table {
   //this is manually added to avoid encoding/decoding cost as our object
   //id is a byte array instead of a string
   public ByteBuffer returnsAsByteBuffer(int j) {
-    int o = __offset(28);
+    int o = __offset(30);
     if (o == 0) {
       return null;
     }

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

@@ -127,6 +127,7 @@ public class RayletClientImpl implements RayletClient {
     UniqueId parentTaskId = UniqueId.fromByteBuffer(info.parentTaskIdAsByteBuffer());
     int parentCounter = info.parentCounter();
     UniqueId actorCreationId = UniqueId.fromByteBuffer(info.actorCreationIdAsByteBuffer());
+    int maxActorReconstructions = info.maxActorReconstructions();
     UniqueId actorId = UniqueId.fromByteBuffer(info.actorIdAsByteBuffer());
     UniqueId actorHandleId = UniqueId.fromByteBuffer(info.actorHandleIdAsByteBuffer());
     int actorCounter = info.actorCounter();
@@ -162,8 +163,9 @@ public class RayletClientImpl implements RayletClient {
     FunctionDescriptor functionDescriptor = new FunctionDescriptor(
         info.functionDescriptor(0), info.functionDescriptor(1), info.functionDescriptor(2)
     );
-    return new TaskSpec(driverId, taskId, parentTaskId, parentCounter, actorCreationId, actorId,
-        actorHandleId, actorCounter, args, returnIds, resources, functionDescriptor);
+    return new TaskSpec(driverId, taskId, parentTaskId, parentCounter, actorCreationId,
+        maxActorReconstructions, actorId, actorHandleId, actorCounter, args, returnIds, resources,
+        functionDescriptor);
   }
 
   private static ByteBuffer convertTaskSpecToFlatbuffer(TaskSpec task) {
@@ -177,10 +179,11 @@ public class RayletClientImpl implements RayletClient {
     final int parentCounter = task.parentCounter;
     final int actorCreateIdOffset = fbb.createString(task.actorCreationId.toByteBuffer());
     final int actorCreateDummyIdOffset = fbb.createString(task.actorId.toByteBuffer());
+    final int maxActorReconstructions = task.maxActorReconstructions;
     final int actorIdOffset = fbb.createString(task.actorId.toByteBuffer());
     final int actorHandleIdOffset = fbb.createString(task.actorHandleId.toByteBuffer());
     final int actorCounter = task.actorCounter;
-    final int functionIdOffset = fbb.createString(UniqueId.NIL.toByteBuffer());
+    final int functionIdOffset = fbb.createString(UniqueId.randomId().toByteBuffer());
     // Serialize args
     int[] argsOffsets = new int[task.args.length];
     for (int i = 0; i < argsOffsets.length; i++) {
@@ -230,13 +233,24 @@ public class RayletClientImpl implements RayletClient {
     int functionDescriptorOffset = fbb.createVectorOfTables(functionDescriptorOffsets);
 
     int root = TaskInfo.createTaskInfo(
-        fbb, driverIdOffset, taskIdOffset,
-        parentTaskIdOffset, parentCounter,
-        actorCreateIdOffset, actorCreateDummyIdOffset,
-        actorIdOffset, actorHandleIdOffset, actorCounter,
-        false, functionIdOffset,
-        argsOffset, returnsOffset, requiredResourcesOffset,
-        requiredPlacementResourcesOffset, Language.JAVA,
+        fbb,
+        driverIdOffset,
+        taskIdOffset,
+        parentTaskIdOffset,
+        parentCounter,
+        actorCreateIdOffset,
+        actorCreateDummyIdOffset,
+        maxActorReconstructions,
+        actorIdOffset,
+        actorHandleIdOffset,
+        actorCounter,
+        false,
+        functionIdOffset,
+        argsOffset,
+        returnsOffset,
+        requiredResourcesOffset,
+        requiredPlacementResourcesOffset,
+        Language.JAVA,
         functionDescriptorOffset);
     fbb.finish(root);
     ByteBuffer buffer = fbb.dataBuffer();

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

@@ -28,6 +28,8 @@ public class TaskSpec {
   // Id for createActor a target actor
   public final UniqueId actorCreationId;
 
+  public final int maxActorReconstructions;
+
   // Actor ID of the task. This is the actor that this task is executed on
   // or NIL_ACTOR_ID if the task is just a normal task.
   public final UniqueId actorId;
@@ -62,14 +64,15 @@ public class TaskSpec {
   }
 
   public TaskSpec(UniqueId driverId, UniqueId taskId, UniqueId parentTaskId, int parentCounter,
-      UniqueId actorCreationId, UniqueId actorId, UniqueId actorHandleId, int actorCounter,
-      FunctionArg[] args, UniqueId[] returnIds,
+      UniqueId actorCreationId, int maxActorReconstructions, UniqueId actorId,
+      UniqueId actorHandleId, int actorCounter, FunctionArg[] args, UniqueId[] returnIds,
       Map<String, Double> resources, FunctionDescriptor functionDescriptor) {
     this.driverId = driverId;
     this.taskId = taskId;
     this.parentTaskId = parentTaskId;
     this.parentCounter = parentCounter;
     this.actorCreationId = actorCreationId;
+    this.maxActorReconstructions = maxActorReconstructions;
     this.actorId = actorId;
     this.actorHandleId = actorHandleId;
     this.actorCounter = actorCounter;