[Java] improve Java API module (#2783)

API module (`ray/java/api` dir) includes all public APIs provided by Ray, it should be the only module that normal Ray users need to face.

The purpose of this PR to first improve the code quality of the API module. Subsequent PRs will improve other modules later. The changes of this PR include the following aspects: 
1) Only keep interfaces in api module, to hide implementation details from users and fix circular dependencies among modules.
2) Document everything in the api module. 
3) Improve naming.
4) Add more tests for API. 
5) Also fix/improve related code in other modules.
6) Remove some unused code.

(Apologize for posting such a large PR. Java worker code has been lack of maintenance for a while. There're a lot of code quality issues that need to be fixed. We plan to use a couple of large PRs to address them. After that, future changes will come in small PRs.)

java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java

@@ -1,18 +1,11 @@
 package org.ray.core.impl;
 
-import java.lang.reflect.Constructor;
-import java.lang.reflect.InvocationTargetException;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.arrow.plasma.ObjectStoreLink;
 import org.apache.arrow.plasma.PlasmaClient;
-import org.ray.api.RayActor;
-import org.ray.api.RayRemote;
-import org.ray.api.UniqueID;
-import org.ray.api.funcs.RayFunc2;
-import org.ray.core.RayRuntime;
-import org.ray.core.UniqueIdHelper;
+import org.ray.core.AbstractRayRuntime;
 import org.ray.core.WorkerContext;
 import org.ray.core.model.RayParameters;
 import org.ray.core.model.WorkerMode;
@@ -29,13 +22,12 @@ import org.ray.spi.impl.NonRayletStateStoreProxyImpl;
 import org.ray.spi.impl.RayletStateStoreProxyImpl;
 import org.ray.spi.impl.RedisClient;
 import org.ray.spi.model.AddressInfo;
-import org.ray.util.exception.TaskExecutionException;
 import org.ray.util.logger.RayLog;
 
 /**
  * native runtime for local box and cluster run.
  */
-public class RayNativeRuntime extends RayRuntime {
+public final class RayNativeRuntime extends AbstractRayRuntime {
 
   static {
     System.err.println("Current working directory is " + System.getProperty("user.dir"));
@@ -46,8 +38,6 @@ public class RayNativeRuntime extends RayRuntime {
   private StateStoreProxy stateStoreProxy;
   private KeyValueStoreLink kvStore = null;
   private RunManager manager = null;
-  private Object actor = null;
-  private UniqueID actorId = UniqueID.NIL;
 
   protected RayNativeRuntime() {
   }
@@ -103,12 +93,12 @@ public class RayNativeRuntime extends RayRuntime {
     if (params.worker_mode != WorkerMode.NONE) {
       String overwrites = "";
       // initialize the links
-      int releaseDelay = RayRuntime.configReader
+      int releaseDelay = AbstractRayRuntime.configReader
           .getIntegerValue("ray", "plasma_default_release_delay", 0,
               "how many release requests should be delayed in plasma client");
 
       if (!params.use_raylet) {
-        ObjectStoreLink plink = new PlasmaClient(params.object_store_name, 
+        ObjectStoreLink plink = new PlasmaClient(params.object_store_name,
             params.object_store_manager_name, releaseDelay);
 
         LocalSchedulerLink slink = new DefaultLocalSchedulerClient(
@@ -152,24 +142,15 @@ public class RayNativeRuntime extends RayRuntime {
   }
 
   @Override
-  public void cleanUp() {
+  public void shutdown() {
     if (null != manager) {
       manager.cleanup(true);
     }
   }
 
-  @Override
-  public Object getLocalActor(UniqueID id) {
-    if (actorId.equals(id)) {
-      return actor;
-    } else {
-      return null;
-    }
-  }
-
   private void startOnebox(RayParameters params, PathConfig paths) throws Exception {
     params.cleanup = true;
-    manager = new RunManager(params, paths, RayRuntime.configReader);
+    manager = new RunManager(params, paths, AbstractRayRuntime.configReader);
     manager.startRayHead();
 
     params.redis_address = manager.info().redisAddress;
@@ -235,59 +216,4 @@ public class RayNativeRuntime extends RayRuntime {
       kvStore.hmset("Workers:" + workerId, workerInfo);
     }
   }
-
-  @SuppressWarnings("unchecked")
-  @Override
-  public <T> RayActor<T> create(Class<T> cls) {
-    UniqueID createTaskId = localSchedulerProxy.generateTaskId(
-        WorkerContext.currentTask().driverId,
-        WorkerContext.currentTask().taskId,
-        WorkerContext.nextCallIndex()
-    );
-
-    UniqueID actorId = UniqueIdHelper.computeReturnId(createTaskId, 0);
-    RayActor<T> actor = new RayActor<>(actorId);
-    UniqueID cursorId;
-
-    RayFunc2<byte[], String, byte[]> createActorLambda = RayNativeRuntime::createActorInActor;
-    cursorId = worker.createActor(
-        createTaskId,
-        actorId,
-        createActorLambda,
-        new Object[]{actorId.getBytes(), cls.getName()}
-    ).getId();
-    actor.setTaskCursor(cursorId);
-    return actor;
-  }
-
-  @RayRemote
-  public static byte[] createActorInActor(byte[] actorId, String className) {
-    ((RayNativeRuntime) RayRuntime.getInstance()).localCreateActorInActor(actorId, className);
-    return actorId;
-  }
-
-  public Object localCreateActorInActor(byte[] actorId, String className) {
-    try {
-      this.actorId = new UniqueID(actorId);
-      Class<?> cls = Class.forName(className, true, Thread.currentThread().getContextClassLoader());
-
-      Constructor<?>[] cts = cls.getConstructors();
-      for (Constructor<?> ct : cts) {
-        System.err.println(ct.getName() + ", param count = " + ct.getParameterCount());
-      }
-
-      actor = cls.getConstructor(new Class<?>[0]).newInstance();
-      RayLog.core.info("create actor " + this.actorId + " inside actor ok");
-      return actor;
-    } catch (ClassNotFoundException | InstantiationException | IllegalAccessException
-        | IllegalArgumentException | InvocationTargetException | NoSuchMethodException
-        | SecurityException e) {
-      e.printStackTrace();
-      String log = "create actor " + this.actorId + " for " + className + "  failed, ex = " + e
-          .getMessage();
-      System.err.println(log);
-      RayLog.core.error(log, e);
-      throw new TaskExecutionException(log, e);
-    }
-  }
 }

java/runtime-native/src/main/java/org/ray/runner/RunManager.java

@@ -12,7 +12,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.TimeUnit;
-import org.ray.api.UniqueID;
+import org.ray.api.id.UniqueId;
 import org.ray.core.model.RayParameters;
 import org.ray.core.model.RunMode;
 import org.ray.runner.RunInfo.ProcessType;
@@ -101,7 +101,7 @@ public class RunManager {
     startRayProcesses();
   }
 
-  public Process startDriver(String mainClass, String redisAddress, UniqueID driverId,
+  public Process startDriver(String mainClass, String redisAddress, UniqueId driverId,
       String logDir, String ip,
       String driverClass, String driverArgs, String additonalClassPaths,
       String additionalConfigs) {
@@ -369,7 +369,7 @@ public class RunManager {
         for (int j = 0; j < localNumWorkers[i]; j++) {
           startWorker(localStores.storeName, localStores.managerName, localStores.schedulerName,
               "/worker" + i + "." + j, params.redis_address,
-              params.node_ip_address, UniqueID.NIL, "", params.redirect, params.cleanup);
+              params.node_ip_address, UniqueId.NIL, "", params.redirect, params.cleanup);
         }
       }
     }
@@ -570,7 +570,7 @@ public class RunManager {
 
     String workerCmd = null;
     workerCmd = buildWorkerCommand(true, info.storeName, info.managerName, name,
-        UniqueID.NIL, "", ip, redisAddress);
+        UniqueId.NIL, "", ip, redisAddress);
     cmd += " -w \"" + workerCmd + "\"";
 
     if (redisAddress.length() > 0) {
@@ -614,7 +614,7 @@ public class RunManager {
 
     //Create the worker command that the raylet will use to start workers.
     String workerCommand = buildWorkerCommandRaylet(info.storeName, rayletSocketName,
-        UniqueID.NIL, "", ip, redisAddress);
+        UniqueId.NIL, "", ip, redisAddress);
 
     int sep = redisAddress.indexOf(':');
     assert (sep != -1);
@@ -654,13 +654,13 @@ public class RunManager {
   }
 
   private String buildWorkerCommandRaylet(String storeName, String rayletSocketName,
-      UniqueID actorId, String actorClass,
+      UniqueId actorId, String actorClass,
       String ip, String redisAddress) {
     String workerConfigs = "ray.java.start.object_store_name=" + storeName
         + ";ray.java.start.raylet_socket_name=" + rayletSocketName
         + ";ray.java.start.worker_mode=WORKER;ray.java.start.use_raylet=true";
     workerConfigs += ";ray.java.start.deploy=" + params.deploy;
-    if (!actorId.equals(UniqueID.NIL)) {
+    if (!actorId.equals(UniqueId.NIL)) {
       workerConfigs += ";ray.java.start.actor_id=" + actorId;
     }
     if (!actorClass.equals("")) {
@@ -685,14 +685,14 @@ public class RunManager {
 
   private String buildWorkerCommand(boolean isFromLocalScheduler, String storeName,
       String storeManagerName, String localSchedulerName,
-      UniqueID actorId, String actorClass, String
+      UniqueId actorId, String actorClass, String
       ip, String redisAddress) {
     String workerConfigs = "ray.java.start.object_store_name=" + storeName
         + ";ray.java.start.object_store_manager_name=" + storeManagerName
         + ";ray.java.start.worker_mode=WORKER"
         + ";ray.java.start.local_scheduler_name=" + localSchedulerName;
     workerConfigs += ";ray.java.start.deploy=" + params.deploy;
-    if (!actorId.equals(UniqueID.NIL)) {
+    if (!actorId.equals(UniqueId.NIL)) {
       workerConfigs += ";ray.java.start.actor_id=" + actorId;
     }
     if (!actorClass.equals("")) {
@@ -783,7 +783,7 @@ public class RunManager {
 
   public void startWorker(String storeName, String storeManagerName,
       String localSchedulerName, String workerName, String redisAddress,
-      String ip, UniqueID actorId, String actorClass,
+      String ip, UniqueId actorId, String actorClass,
       boolean redirect, boolean cleanup) {
     String cmd = buildWorkerCommand(false, storeName, storeManagerName, localSchedulerName, actorId,
         actorClass, ip, redisAddress);

java/runtime-native/src/main/java/org/ray/runner/worker/DefaultDriver.java

@@ -1,6 +1,6 @@
 package org.ray.runner.worker;
 
-import org.ray.core.RayRuntime;
+import org.ray.core.AbstractRayRuntime;
 import org.ray.core.model.WorkerMode;
 
 /**
@@ -15,13 +15,13 @@ public class DefaultDriver {
   //
   public static void main(String[] args) {
     try {
-      RayRuntime.init(args);
-      assert RayRuntime.getParams().worker_mode == WorkerMode.DRIVER;
+      AbstractRayRuntime.init(args);
+      assert AbstractRayRuntime.getParams().worker_mode == WorkerMode.DRIVER;
 
-      String driverClass = RayRuntime.configReader
+      String driverClass = AbstractRayRuntime.configReader
           .getStringValue("ray.java.start", "driver_class", "",
               "java class which main is served as the driver in a java worker");
-      String driverArgs = RayRuntime.configReader
+      String driverArgs = AbstractRayRuntime.configReader
           .getStringValue("ray.java.start", "driver_args", "",
               "arguments for the java class main function which is served at the driver");
       Class<?> cls = Class.forName(driverClass);

java/runtime-native/src/main/java/org/ray/runner/worker/DefaultWorker.java

@@ -1,6 +1,6 @@
 package org.ray.runner.worker;
 
-import org.ray.core.RayRuntime;
+import org.ray.core.AbstractRayRuntime;
 import org.ray.core.model.WorkerMode;
 
 /**
@@ -16,9 +16,9 @@ public class DefaultWorker {
   //
   public static void main(String[] args) {
     try {
-      RayRuntime.init(args);
-      assert RayRuntime.getParams().worker_mode == WorkerMode.WORKER;
-      RayRuntime.getInstance().loop();
+      AbstractRayRuntime.init(args);
+      assert AbstractRayRuntime.getParams().worker_mode == WorkerMode.WORKER;
+      AbstractRayRuntime.getInstance().loop();
       throw new RuntimeException("Control flow should never reach here");
 
     } catch (Throwable e) {

java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java

@@ -7,9 +7,8 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
-import org.ray.api.Ray;
-import org.ray.api.UniqueID;
-import org.ray.core.RayRuntime;
+import org.ray.api.id.UniqueId;
+import org.ray.core.AbstractRayRuntime;
 import org.ray.core.UniqueIdHelper;
 import org.ray.spi.LocalSchedulerLink;
 import org.ray.spi.model.FunctionArg;
@@ -24,15 +23,15 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
 
   private static ThreadLocal<ByteBuffer> _taskBuffer = ThreadLocal.withInitial(() -> {
     ByteBuffer bb = ByteBuffer
-        .allocateDirect(RayRuntime.getParams().max_submit_task_buffer_size_bytes);
+        .allocateDirect(AbstractRayRuntime.getParams().max_submit_task_buffer_size_bytes);
     bb.order(ByteOrder.LITTLE_ENDIAN);
     return bb;
   });
   private long client = 0;
   boolean useRaylet = false;
 
-  public DefaultLocalSchedulerClient(String schedulerSockName, UniqueID clientId,
-      boolean isWorker, UniqueID driverId, boolean useRaylet) {
+  public DefaultLocalSchedulerClient(String schedulerSockName, UniqueId clientId,
+      boolean isWorker, UniqueId driverId, boolean useRaylet) {
     client = nativeInit(schedulerSockName, clientId.getBytes(),
         isWorker, driverId.getBytes(), useRaylet);
     this.useRaylet = useRaylet;
@@ -57,6 +56,7 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
 
   @Override
   public void submitTask(TaskSpec task) {
+    RayLog.core.debug("Submitting task: {}", task);
     // We don't support resources management in non raylet mode.
     if (!useRaylet) {
       task.resources.clear();
@@ -72,12 +72,12 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     }
 
     ByteBuffer info = taskSpec2Info(task);
-    byte[] a = null;
+    byte[] cursorId = null;
     if (!task.actorId.isNil()) {
-      a = task.cursorId.getBytes();
+      cursorId = task.cursorId.getBytes();
     }
 
-    nativeSubmitTask(client, a, info, info.position(), info.remaining(), useRaylet);
+    nativeSubmitTask(client, cursorId, info, info.position(), info.remaining(), useRaylet);
   }
 
   @Override
@@ -89,19 +89,19 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
   }
 
   @Override
-  public void markTaskPutDependency(UniqueID taskId, UniqueID objectId) {
+  public void markTaskPutDependency(UniqueId taskId, UniqueId objectId) {
     nativePutObject(client, taskId.getBytes(), objectId.getBytes());
   }
 
   @Override
-  public void reconstructObject(UniqueID objectId, boolean fetchOnly) {
-    List<UniqueID> objects = new ArrayList<>();
+  public void reconstructObject(UniqueId objectId, boolean fetchOnly) {
+    List<UniqueId> objects = new ArrayList<>();
     objects.add(objectId);
     nativeReconstructObjects(client, getIdBytes(objects), fetchOnly);
   }
 
   @Override
-  public void reconstructObjects(List<UniqueID> objectIds, boolean fetchOnly) {
+  public void reconstructObjects(List<UniqueId> objectIds, boolean fetchOnly) {
     if (RayLog.core.isInfoEnabled()) {
       RayLog.core.info("Reconstructing objects for task {}, object IDs are {}",
           UniqueIdHelper.computeTaskId(objectIds.get(0)), objectIds);
@@ -110,9 +110,9 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
   }
 
   @Override
-  public UniqueID generateTaskId(UniqueID driverId, UniqueID parentTaskId, int taskIndex) {
+  public UniqueId generateTaskId(UniqueId driverId, UniqueId parentTaskId, int taskIndex) {
     byte[] bytes = nativeGenerateTaskId(driverId.getBytes(), parentTaskId.getBytes(), taskIndex);
-    return new UniqueID(bytes);
+    return new UniqueId(bytes);
   }
 
   @Override
@@ -125,49 +125,47 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     TaskInfo info = TaskInfo.getRootAsTaskInfo(bb);
 
     TaskSpec spec = new TaskSpec();
-    spec.driverId = UniqueID.fromByteBuffer(info.driverIdAsByteBuffer());
-    spec.taskId = UniqueID.fromByteBuffer(info.taskIdAsByteBuffer());
-    spec.parentTaskId = UniqueID.fromByteBuffer(info.parentTaskIdAsByteBuffer());
+    spec.driverId = UniqueId.fromByteBuffer(info.driverIdAsByteBuffer());
+    spec.taskId = UniqueId.fromByteBuffer(info.taskIdAsByteBuffer());
+    spec.parentTaskId = UniqueId.fromByteBuffer(info.parentTaskIdAsByteBuffer());
     spec.parentCounter = info.parentCounter();
-    spec.actorId = UniqueID.fromByteBuffer(info.actorIdAsByteBuffer());
+    spec.actorId = UniqueId.fromByteBuffer(info.actorIdAsByteBuffer());
     spec.actorCounter = info.actorCounter();
-    spec.createActorId = UniqueID.fromByteBuffer(info.actorCreationIdAsByteBuffer());
+    spec.createActorId = UniqueId.fromByteBuffer(info.actorCreationIdAsByteBuffer());
 
-    spec.functionId = UniqueID.fromByteBuffer(info.functionIdAsByteBuffer());
+    spec.functionId = UniqueId.fromByteBuffer(info.functionIdAsByteBuffer());
 
     List<FunctionArg> args = new ArrayList<>();
     for (int i = 0; i < info.argsLength(); i++) {
-      FunctionArg darg = new FunctionArg();
+      UniqueId id = null;
+      byte[] data = null;
       Arg sarg = info.args(i);
 
       int idCount = sarg.objectIdsLength();
       if (idCount > 0) {
-        darg.ids = new ArrayList<>();
-        for (int j = 0; j < idCount; j++) {
-          ByteBuffer lbb = sarg.objectIdAsByteBuffer(j);
-          assert (lbb != null && lbb.remaining() > 0);
-          darg.ids.add(UniqueID.fromByteBuffer(lbb));
-        }
+        ByteBuffer lbb = sarg.objectIdAsByteBuffer(0);
+        assert (lbb != null && lbb.remaining() > 0);
+        id = UniqueId.fromByteBuffer(lbb);
       }
 
       ByteBuffer lbb = sarg.dataAsByteBuffer();
       if (lbb != null && lbb.remaining() > 0) {
         // TODO: how to avoid memory copy
-        darg.data = new byte[lbb.remaining()];
-        lbb.get(darg.data);
+        data = new byte[lbb.remaining()];
+        lbb.get(data);
       }
 
-      args.add(darg);
+      args.add(new FunctionArg(id, data));
     }
     spec.args = args.toArray(new FunctionArg[0]);
 
-    List<UniqueID> rids = new ArrayList<>();
+    List<UniqueId> rids = new ArrayList<>();
     for (int i = 0; i < info.returnsLength(); i++) {
       ByteBuffer lbb = info.returnsAsByteBuffer(i);
       assert (lbb != null && lbb.remaining() > 0);
-      rids.add(UniqueID.fromByteBuffer(lbb));
+      rids.add(UniqueId.fromByteBuffer(lbb));
     }
-    spec.returnIds = rids.toArray(new UniqueID[0]);
+    spec.returnIds = rids.toArray(new UniqueId[0]);
 
     return spec;
   }
@@ -183,7 +181,7 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     final int parentTaskIdOffset = fbb.createString(task.parentTaskId.toByteBuffer());
     final int parentCounter = task.parentCounter;
     final int actorCreateIdOffset = fbb.createString(task.createActorId.toByteBuffer());
-    final int actorCreateDummyIdOffset = fbb.createString(UniqueID.NIL.toByteBuffer());
+    final int actorCreateDummyIdOffset = fbb.createString(UniqueId.NIL.toByteBuffer());
     final int actorIdOffset = fbb.createString(task.actorId.toByteBuffer());
     final int actorHandleIdOffset = fbb.createString(task.actorHandleId.toByteBuffer());
     final int actorCounter = task.actorCounter;
@@ -195,12 +193,10 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
 
       int objectIdOffset = 0;
       int dataOffset = 0;
-      if (task.args[i].ids != null) {
-        int idCount = task.args[i].ids.size();
-        int[] idOffsets = new int[idCount];
-        for (int k = 0; k < idCount; k++) {
-          idOffsets[k] = fbb.createString(task.args[i].ids.get(k).toByteBuffer());
-        }
+      if (task.args[i].id != null) {
+        int[] idOffsets = new int[] {
+            fbb.createString(task.args[i].id.toByteBuffer())
+        };
         objectIdOffset = fbb.createVectorOfTables(idOffsets);
       } else {
         objectIdOffset = fbb.createVectorOfTables(new int[0]);
@@ -248,9 +244,9 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     fbb.finish(root);
     ByteBuffer buffer = fbb.dataBuffer();
 
-    if (buffer.remaining() > RayRuntime.getParams().max_submit_task_buffer_size_bytes) {
+    if (buffer.remaining() > AbstractRayRuntime.getParams().max_submit_task_buffer_size_bytes) {
       RayLog.core.error(
-          "Allocated buffer is not enough to transfer the task specification: " + RayRuntime
+          "Allocated buffer is not enough to transfer the task specification: " + AbstractRayRuntime
               .getParams().max_submit_task_buffer_size_bytes + " vs " + buffer.remaining());
       assert (false);
     }
@@ -258,7 +254,7 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     return buffer;
   }
 
-  private static byte[][] getIdBytes(List<UniqueID> objectIds) {
+  private static byte[][] getIdBytes(List<UniqueId> objectIds) {
     int size = objectIds.size();
     byte[][] ids = new byte[size][];
     for (int i = 0; i < size; i++) {

java/runtime-native/src/main/java/org/ray/spi/impl/NativeRemoteFunctionManager.java

@@ -5,7 +5,7 @@ import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.concurrent.ConcurrentHashMap;
 import net.lingala.zip4j.core.ZipFile;
-import org.ray.api.UniqueID;
+import org.ray.api.id.UniqueId;
 import org.ray.spi.KeyValueStoreLink;
 import org.ray.spi.RemoteFunctionManager;
 import org.ray.util.FileUtil;
@@ -18,7 +18,7 @@ import org.ray.util.logger.RayLog;
  */
 public class NativeRemoteFunctionManager implements RemoteFunctionManager {
 
-  private final ConcurrentHashMap<UniqueID, ClassLoader> loadedApps = new ConcurrentHashMap<>();
+  private final ConcurrentHashMap<UniqueId, ClassLoader> loadedApps = new ConcurrentHashMap<>();
   private MessageDigest md;
   private final String appDir = System.getProperty("user.dir") + "/apps";
   private final KeyValueStoreLink kvStore;
@@ -35,11 +35,11 @@ public class NativeRemoteFunctionManager implements RemoteFunctionManager {
   }
 
   @Override
-  public UniqueID registerResource(byte[] resourceZip) {
+  public UniqueId registerResource(byte[] resourceZip) {
     byte[] digest = Sha1Digestor.digest(resourceZip);
-    assert (digest.length == UniqueID.LENGTH);
+    assert (digest.length == UniqueId.LENGTH);
 
-    UniqueID resourceId = new UniqueID(digest);
+    UniqueId resourceId = new UniqueId(digest);
 
     // TODO: resources must be saved in persistent store
     kvStore.set(resourceId.getBytes(), resourceZip, null);
@@ -48,32 +48,32 @@ public class NativeRemoteFunctionManager implements RemoteFunctionManager {
   }
 
   @Override
-  public byte[] getResource(UniqueID resourceId) {
+  public byte[] getResource(UniqueId resourceId) {
     return kvStore.get(resourceId.getBytes(), null);
   }
 
   @Override
-  public void unregisterResource(UniqueID resourceId) {
+  public void unregisterResource(UniqueId resourceId) {
     kvStore.delete(resourceId.getBytes(), null);
   }
 
   @Override
-  public void registerApp(UniqueID driverId, UniqueID resourceId) {
+  public void registerApp(UniqueId driverId, UniqueId resourceId) {
     kvStore.set("App2ResMap", resourceId.toString(), driverId.toString());
   }
 
   @Override
-  public UniqueID getAppResourceId(UniqueID driverId) {
-    return UniqueID.fromHexString(kvStore.get("App2ResMap", driverId.toString()));
+  public UniqueId getAppResourceId(UniqueId driverId) {
+    return UniqueId.fromHexString(kvStore.get("App2ResMap", driverId.toString()));
   }
 
   @Override
-  public void unregisterApp(UniqueID driverId) {
+  public void unregisterApp(UniqueId driverId) {
     kvStore.delete("App2ResMap", driverId.toString());
   }
 
   @Override
-  public ClassLoader loadResource(UniqueID driverId) {
+  public ClassLoader loadResource(UniqueId driverId) {
     ClassLoader classLoader = loadedApps.get(driverId);
     if (classLoader == null) {
       synchronized (this) {
@@ -86,13 +86,13 @@ public class NativeRemoteFunctionManager implements RemoteFunctionManager {
     return classLoader;
   }
 
-  private ClassLoader initLoadedApps(UniqueID driverId) {
+  private ClassLoader initLoadedApps(UniqueId driverId) {
     try {
       RayLog.core.info("initLoadedApps" + driverId.toString());
 
       ClassLoader cl = loadedApps.get(driverId);
       if (cl == null) {
-        UniqueID resId = UniqueID.fromHexString(kvStore.get("App2ResMap", driverId.toString()));
+        UniqueId resId = UniqueId.fromHexString(kvStore.get("App2ResMap", driverId.toString()));
         byte[] res = getResource(resId);
         if (res == null) {
           throw new RuntimeException("get resource null, the resId " + resId.toString());
@@ -120,7 +120,7 @@ public class NativeRemoteFunctionManager implements RemoteFunctionManager {
   }
 
   @Override
-  public synchronized void unloadFunctions(UniqueID driverId) {
+  public synchronized void unloadFunctions(UniqueId driverId) {
     ClassLoader cl = loadedApps.get(driverId);
     try {
       JarLoader.unloadJars(cl);

java/runtime-native/src/main/java/org/ray/spi/impl/RayletStateStoreProxyImpl.java

@@ -5,7 +5,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
-import org.ray.api.UniqueID;
+import org.ray.api.id.UniqueId;
 import org.ray.format.gcs.ClientTableData;
 import org.ray.spi.KeyValueStoreLink;
 import org.ray.spi.model.AddressInfo;
@@ -29,7 +29,7 @@ public class RayletStateStoreProxyImpl extends BaseStateStoreProxyImpl {
     List<AddressInfo> schedulerInfo = new ArrayList<>();
 
     byte[] prefix = "CLIENT".getBytes();
-    byte[] postfix = UniqueID.genNil().getBytes();
+    byte[] postfix = UniqueId.genNil().getBytes();
     byte[] clientKey = new byte[prefix.length + postfix.length];
     System.arraycopy(prefix, 0, clientKey, 0, prefix.length);
     System.arraycopy(postfix, 0, clientKey, prefix.length, postfix.length);