[Java] Support direct actor call in Java worker (#5504)

java/api/src/main/java/org/ray/api/options/ActorCreationOptions.java

@@ -10,26 +10,33 @@ public class ActorCreationOptions extends BaseTaskOptions {
 
   public static final int NO_RECONSTRUCTION = 0;
   public static final int INFINITE_RECONSTRUCTIONS = (int) Math.pow(2, 30);
+  // DO NOT set this environment variable. It's only used for test purposes.
+  // Please use `setUseDirectCall` instead.
+  public static final boolean DEFAULT_USE_DIRECT_CALL = "1"
+      .equals(System.getenv("ACTOR_CREATION_OPTIONS_DEFAULT_USE_DIRECT_CALL"));
 
   public final int maxReconstructions;
 
+  public final boolean useDirectCall;
+
   public final String jvmOptions;
 
-  private ActorCreationOptions(Map<String, Double> resources,
-                               int maxReconstructions,
-                               String jvmOptions) {
+  private ActorCreationOptions(Map<String, Double> resources, int maxReconstructions,
+      boolean useDirectCall, String jvmOptions) {
     super(resources);
     this.maxReconstructions = maxReconstructions;
+    this.useDirectCall = useDirectCall;
     this.jvmOptions = jvmOptions;
   }
 
   /**
-   *  The inner class for building ActorCreationOptions.
+   * The inner class for building ActorCreationOptions.
    */
   public static class Builder {
 
     private Map<String, Double> resources = new HashMap<>();
     private int maxReconstructions = NO_RECONSTRUCTION;
+    private boolean useDirectCall = DEFAULT_USE_DIRECT_CALL;
     private String jvmOptions = null;
 
     public Builder setResources(Map<String, Double> resources) {
@@ -42,13 +49,21 @@ public Builder setMaxReconstructions(int maxReconstructions) {
       return this;
     }
 
+    // Since direct call is not fully supported yet (see issue #5559),
+    // users are not allowed to set the option to true.
+    // TODO (kfstorm): uncomment when direct call is ready.
+//    public Builder setUseDirectCall(boolean useDirectCall) {
+//      this.useDirectCall = useDirectCall;
+//      return this;
+//    }
+
     public Builder setJvmOptions(String jvmOptions) {
       this.jvmOptions = jvmOptions;
       return this;
     }
 
     public ActorCreationOptions createActorCreationOptions() {
-      return new ActorCreationOptions(resources, maxReconstructions, jvmOptions);
+      return new ActorCreationOptions(resources, maxReconstructions, useDirectCall, jvmOptions);
     }
   }
 

java/dependencies.bzl

@@ -16,7 +16,7 @@ def gen_java_deps():
             "org.apache.commons:commons-lang3:3.4",
             "org.ow2.asm:asm:6.0",
             "org.slf4j:slf4j-log4j12:1.7.25",
-            "org.testng:testng:6.9.9",
+            "org.testng:testng:6.9.10",
             "redis.clients:jedis:2.8.0",
         ],
         repositories = [

java/runtime/pom.xml

@@ -75,7 +75,7 @@
 <dependency>
   <groupId>org.testng</groupId>
   <artifactId>testng</artifactId>
-  <version>6.9.9</version>
+  <version>6.9.10</version>
 </dependency>
 <dependency>
   <groupId>redis.clients</groupId>

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

@@ -13,11 +13,11 @@
 import org.ray.api.function.RayFunc;
 import org.ray.api.function.RayFuncVoid;
 import org.ray.api.id.ObjectId;
-import org.ray.api.id.UniqueId;
 import org.ray.api.options.ActorCreationOptions;
 import org.ray.api.options.CallOptions;
 import org.ray.api.runtime.RayRuntime;
 import org.ray.api.runtimecontext.RuntimeContext;
+import org.ray.runtime.actor.NativeRayActor;
 import org.ray.runtime.config.RayConfig;
 import org.ray.runtime.context.RuntimeContextImpl;
 import org.ray.runtime.context.WorkerContext;
@@ -28,7 +28,6 @@
 import org.ray.runtime.generated.Common.Language;
 import org.ray.runtime.object.ObjectStore;
 import org.ray.runtime.object.RayObjectImpl;
-import org.ray.runtime.raylet.RayletClient;
 import org.ray.runtime.task.ArgumentsBuilder;
 import org.ray.runtime.task.FunctionArg;
 import org.ray.runtime.task.TaskExecutor;
@@ -51,7 +50,6 @@ public abstract class AbstractRayRuntime implements RayRuntime {
 
   protected ObjectStore objectStore;
   protected TaskSubmitter taskSubmitter;
-  protected RayletClient rayletClient;
   protected WorkerContext workerContext;
 
   public AbstractRayRuntime(RayConfig rayConfig) {
@@ -85,15 +83,6 @@ public void free(List<ObjectId> objectIds, boolean localOnly, boolean deleteCrea
     objectStore.delete(objectIds, localOnly, deleteCreatingTasks);
   }
 
-  @Override
-  public void setResource(String resourceName, double capacity, UniqueId nodeId) {
-    Preconditions.checkArgument(Double.compare(capacity, 0) >= 0);
-    if (nodeId == null) {
-      nodeId = UniqueId.NIL;
-    }
-    rayletClient.setResource(resourceName, capacity, nodeId);
-  }
-
   @Override
   public <T> WaitResult<T> wait(List<RayObject<T>> waitList, int numReturns, int timeoutMs) {
     return objectStore.wait(waitList, numReturns, timeoutMs);
@@ -176,7 +165,7 @@ public Callable wrapCallable(Callable callable) {
 
   private RayObject callNormalFunction(FunctionDescriptor functionDescriptor,
       Object[] args, int numReturns, CallOptions options) {
-    List<FunctionArg> functionArgs = ArgumentsBuilder.wrap(args);
+    List<FunctionArg> functionArgs = ArgumentsBuilder.wrap(args, /*isDirectCall*/false);
     List<ObjectId> returnIds = taskSubmitter.submitTask(functionDescriptor,
         functionArgs, numReturns, options);
     Preconditions.checkState(returnIds.size() == numReturns && returnIds.size() <= 1);
@@ -189,7 +178,7 @@ private RayObject callNormalFunction(FunctionDescriptor functionDescriptor,
 
   private RayObject callActorFunction(RayActor rayActor,
       FunctionDescriptor functionDescriptor, Object[] args, int numReturns) {
-    List<FunctionArg> functionArgs = ArgumentsBuilder.wrap(args);
+    List<FunctionArg> functionArgs = ArgumentsBuilder.wrap(args, isDirectCall(rayActor));
     List<ObjectId> returnIds = taskSubmitter.submitActorTask(rayActor,
         functionDescriptor, functionArgs, numReturns, null);
     Preconditions.checkState(returnIds.size() == numReturns && returnIds.size() <= 1);
@@ -202,14 +191,21 @@ private RayObject callActorFunction(RayActor rayActor,
 
   private RayActor createActorImpl(FunctionDescriptor functionDescriptor,
       Object[] args, ActorCreationOptions options) {
-    List<FunctionArg> functionArgs = ArgumentsBuilder.wrap(args);
+    List<FunctionArg> functionArgs = ArgumentsBuilder.wrap(args, /*isDirectCall*/false);
     if (functionDescriptor.getLanguage() != Language.JAVA && options != null) {
       Preconditions.checkState(Strings.isNullOrEmpty(options.jvmOptions));
     }
     RayActor actor = taskSubmitter.createActor(functionDescriptor, functionArgs, options);
     return actor;
   }
 
+  private boolean isDirectCall(RayActor rayActor) {
+    if (rayActor instanceof NativeRayActor) {
+      return ((NativeRayActor) rayActor).isDirectCallActor();
+    }
+    return false;
+  }
+
   public WorkerContext getWorkerContext() {
     return workerContext;
   }
@@ -218,10 +214,6 @@ public ObjectStore getObjectStore() {
     return objectStore;
   }
 
-  public RayletClient getRayletClient() {
-    return rayletClient;
-  }
-
   public FunctionManager getFunctionManager() {
     return functionManager;
   }

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

@@ -2,37 +2,45 @@
 
 import java.util.concurrent.atomic.AtomicInteger;
 import org.ray.api.id.JobId;
+import org.ray.api.id.UniqueId;
 import org.ray.runtime.config.RayConfig;
 import org.ray.runtime.context.LocalModeWorkerContext;
 import org.ray.runtime.object.LocalModeObjectStore;
-import org.ray.runtime.raylet.LocalModeRayletClient;
+import org.ray.runtime.task.LocalModeTaskExecutor;
 import org.ray.runtime.task.LocalModeTaskSubmitter;
-import org.ray.runtime.task.TaskExecutor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class RayDevRuntime extends AbstractRayRuntime {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(RayDevRuntime.class);
+
   private AtomicInteger jobCounter = new AtomicInteger(0);
 
   public RayDevRuntime(RayConfig rayConfig) {
     super(rayConfig);
     if (rayConfig.getJobId().isNil()) {
       rayConfig.setJobId(nextJobId());
     }
-    taskExecutor = new TaskExecutor(this);
+    taskExecutor = new LocalModeTaskExecutor(this);
     workerContext = new LocalModeWorkerContext(rayConfig.getJobId());
     objectStore = new LocalModeObjectStore(workerContext);
     taskSubmitter = new LocalModeTaskSubmitter(this, (LocalModeObjectStore) objectStore,
         rayConfig.numberExecThreadsForDevRuntime);
     ((LocalModeObjectStore) objectStore).addObjectPutCallback(
         objectId -> ((LocalModeTaskSubmitter) taskSubmitter).onObjectPut(objectId));
-    rayletClient = new LocalModeRayletClient();
   }
 
   @Override
   public void shutdown() {
     taskExecutor = null;
   }
 
+  @Override
+  public void setResource(String resourceName, double capacity, UniqueId nodeId) {
+    LOGGER.error("Not implemented under SINGLE_PROCESS mode.");
+  }
+
   private JobId nextJobId() {
     return JobId.fromInt(jobCounter.getAndIncrement());
   }

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

@@ -9,15 +9,16 @@
 import java.util.Map;
 import org.apache.commons.io.FileUtils;
 import org.ray.api.id.JobId;
+import org.ray.api.id.UniqueId;
 import org.ray.runtime.config.RayConfig;
 import org.ray.runtime.context.NativeWorkerContext;
 import org.ray.runtime.gcs.GcsClient;
 import org.ray.runtime.gcs.GcsClientOptions;
 import org.ray.runtime.gcs.RedisClient;
 import org.ray.runtime.generated.Common.WorkerType;
 import org.ray.runtime.object.NativeObjectStore;
-import org.ray.runtime.raylet.NativeRayletClient;
 import org.ray.runtime.runner.RunManager;
+import org.ray.runtime.task.NativeTaskExecutor;
 import org.ray.runtime.task.NativeTaskSubmitter;
 import org.ray.runtime.task.TaskExecutor;
 import org.ray.runtime.util.FileUtil;
@@ -112,11 +113,10 @@ public RayNativeRuntime(RayConfig rayConfig) {
         new GcsClientOptions(rayConfig));
     Preconditions.checkState(nativeCoreWorkerPointer != 0);
 
-    taskExecutor = new TaskExecutor(this);
+    taskExecutor = new NativeTaskExecutor(nativeCoreWorkerPointer, this);
     workerContext = new NativeWorkerContext(nativeCoreWorkerPointer);
     objectStore = new NativeObjectStore(workerContext, nativeCoreWorkerPointer);
     taskSubmitter = new NativeTaskSubmitter(nativeCoreWorkerPointer);
-    rayletClient = new NativeRayletClient(nativeCoreWorkerPointer);
 
     // register
     registerWorker();
@@ -136,6 +136,15 @@ public void shutdown() {
     }
   }
 
+  @Override
+  public void setResource(String resourceName, double capacity, UniqueId nodeId) {
+    Preconditions.checkArgument(Double.compare(capacity, 0) >= 0);
+    if (nodeId == null) {
+      nodeId = UniqueId.NIL;
+    }
+    nativeSetResource(nativeCoreWorkerPointer, resourceName, capacity, nodeId.getBytes());
+  }
+
   public void run() {
     nativeRunTaskExecutor(nativeCoreWorkerPointer, taskExecutor);
   }
@@ -176,4 +185,7 @@ private static native void nativeRunTaskExecutor(long nativeCoreWorkerPointer,
   private static native void nativeSetup(String logDir);
 
   private static native void nativeShutdownHook();
+
+  private static native void nativeSetResource(long conn, String resourceName, double capacity,
+      byte[] nodeId);
 }

java/runtime/src/main/java/org/ray/runtime/actor/NativeRayActor.java

@@ -51,6 +51,10 @@ public Language getLanguage() {
     return Language.forNumber(nativeGetLanguage(nativeActorHandle));
   }
 
+  public boolean isDirectCallActor() {
+    return nativeIsDirectCallActor(nativeActorHandle);
+  }
+
   @Override
   public String getModuleName() {
     Preconditions.checkState(getLanguage() == Language.PYTHON);
@@ -90,6 +94,8 @@ protected void finalize() {
 
   private static native int nativeGetLanguage(long nativeActorHandle);
 
+  private static native boolean nativeIsDirectCallActor(long nativeActorHandle);
+
   private static native List<String> nativeGetActorCreationTaskFunctionDescriptor(
       long nativeActorHandle);