[Java] Add runtime context (#4194)

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

@@ -10,6 +10,7 @@ import java.util.Map;
 import java.util.stream.Collectors;
 import org.ray.api.RayActor;
 import org.ray.api.RayObject;
+import org.ray.api.RuntimeContext;
 import org.ray.api.WaitResult;
 import org.ray.api.exception.RayException;
 import org.ray.api.function.RayFunc;
@@ -61,6 +62,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
   protected RayletClient rayletClient;
   protected ObjectStoreProxy objectStoreProxy;
   protected FunctionManager functionManager;
+  protected RuntimeContext runtimeContext;
 
   public AbstractRayRuntime(RayConfig rayConfig) {
     this.rayConfig = rayConfig;
@@ -68,6 +70,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
     worker = new Worker(this);
     workerContext = new WorkerContext(rayConfig.workerMode,
             rayConfig.driverId, rayConfig.runMode);
+    runtimeContext = new RuntimeContextImpl(this);
   }
 
   /**
@@ -346,4 +349,9 @@ public abstract class AbstractRayRuntime implements RayRuntime {
   public RayConfig getRayConfig() {
     return rayConfig;
   }
+
+  public RuntimeContext getRuntimeContext() {
+    return runtimeContext;
+  }
+
 }

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

@@ -173,4 +173,22 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
     checkpoints.sort((x, y) -> Long.compare(y.timestamp, x.timestamp));
     return checkpoints;
   }
+
+
+  /**
+   * Query whether the actor exists in Gcs.
+   */
+  boolean actorExistsInGcs(UniqueId actorId) {
+    byte[] key = ArrayUtils.addAll("ACTOR".getBytes(), actorId.getBytes());
+
+    // TODO(qwang): refactor this with `GlobalState` after this issue
+    // getting finished. https://github.com/ray-project/ray/issues/3933
+    for (RedisClient client : redisClients) {
+      if (client.exists(key)) {
+        return true;
+      }
+    }
+
+    return false;
+  }
 }

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

@@ -0,0 +1,56 @@
+package org.ray.runtime;
+
+import com.google.common.base.Preconditions;
+import org.ray.api.RuntimeContext;
+import org.ray.api.id.UniqueId;
+import org.ray.runtime.config.RunMode;
+import org.ray.runtime.config.WorkerMode;
+import org.ray.runtime.task.TaskSpec;
+
+public class RuntimeContextImpl implements RuntimeContext {
+
+  private AbstractRayRuntime runtime;
+
+  public RuntimeContextImpl(AbstractRayRuntime runtime) {
+    this.runtime = runtime;
+  }
+
+  @Override
+  public UniqueId getCurrentDriverId() {
+    return runtime.getWorkerContext().getCurrentDriverId();
+  }
+
+  @Override
+  public UniqueId getCurrentActorId() {
+    Preconditions.checkState(runtime.rayConfig.workerMode == WorkerMode.WORKER);
+    return runtime.getWorker().getCurrentActorId();
+  }
+
+  @Override
+  public boolean wasCurrentActorReconstructed() {
+    TaskSpec currentTask = runtime.getWorkerContext().getCurrentTask();
+    Preconditions.checkState(currentTask != null && currentTask.isActorCreationTask(),
+        "This method can only be called from an actor creation task.");
+    if (isSingleProcess()) {
+      return false;
+    }
+
+    return ((RayNativeRuntime) runtime).actorExistsInGcs(getCurrentActorId());
+  }
+
+  @Override
+  public String getRayletSocketName() {
+    return runtime.getRayConfig().rayletSocketName;
+  }
+
+  @Override
+  public String getObjectStoreSocketName() {
+    return runtime.getRayConfig().objectStoreSocketName;
+  }
+
+  @Override
+  public boolean isSingleProcess() {
+    return RunMode.SINGLE_PROCESS == runtime.getRayConfig().runMode;
+  }
+
+}

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

@@ -63,6 +63,10 @@ public class Worker {
     this.runtime = runtime;
   }
 
+  public UniqueId getCurrentActorId() {
+    return currentActorId;
+  }
+
   public void loop() {
     while (true) {
       LOGGER.info("Fetching new task in thread {}.", Thread.currentThread().getName());
@@ -86,6 +90,11 @@ public class Worker {
       // Set context
       runtime.getWorkerContext().setCurrentTask(spec, rayFunction.classLoader);
       Thread.currentThread().setContextClassLoader(rayFunction.classLoader);
+
+      if (spec.isActorCreationTask()) {
+        currentActorId = returnId;
+      }
+
       // Get local actor object and arguments.
       Object actor = null;
       if (spec.isActorTask()) {
@@ -94,6 +103,7 @@ public class Worker {
           throw actorCreationException;
         }
         actor = currentActor;
+
       }
       Object[] args = ArgumentsBuilder.unwrap(spec, rayFunction.classLoader);
       // Execute the task.
@@ -112,7 +122,6 @@ public class Worker {
       } else {
         maybeLoadCheckpoint(result, returnId);
         currentActor = result;
-        currentActorId = returnId;
       }
       LOGGER.info("Finished executing task {}", spec.taskId);
     } catch (Exception e) {
@@ -121,7 +130,6 @@ public class Worker {
         runtime.put(returnId, new RayTaskException("Error executing task " + spec, e));
       } else {
         actorCreationException = e;
-        currentActorId = returnId;
       }
     } finally {
       Thread.currentThread().setContextClassLoader(oldLoader);

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

@@ -26,6 +26,8 @@ public class WorkerContext {
    */
   private ThreadLocal<Integer> taskIndex;
 
+  private ThreadLocal<TaskSpec> currentTask;
+
   private UniqueId currentDriverId;
 
   private ClassLoader currentClassLoader;
@@ -46,6 +48,7 @@ public class WorkerContext {
     putIndex = ThreadLocal.withInitial(() -> 0);
     currentTaskId = ThreadLocal.withInitial(UniqueId::randomId);
     this.runMode = runMode;
+    currentTask = ThreadLocal.withInitial(() -> null);
     currentClassLoader = null;
     if (workerMode == WorkerMode.DRIVER) {
       workerId = driverId;
@@ -83,6 +86,7 @@ public class WorkerContext {
     this.currentDriverId = task.driverId;
     taskIndex.set(0);
     putIndex.set(0);
+    this.currentTask.set(task);
     currentClassLoader = classLoader;
   }
 
@@ -124,4 +128,10 @@ public class WorkerContext {
     return currentClassLoader;
   }
 
+  /**
+   * Get the current task.
+   */
+  public TaskSpec getCurrentTask() {
+    return this.currentTask.get();
+  }
 }

java/runtime/src/main/java/org/ray/runtime/gcs/RedisClient.java

@@ -85,4 +85,14 @@ public class RedisClient {
       return jedis.lrange(key, start, end);
     }
   }
+
+  /**
+   * Whether the key exists in Redis.
+   */
+  public boolean exists(byte[] key) {
+    try (Jedis jedis = jedisPool.getResource()) {
+      return jedis.exists(key);
+    }
+  }
+
 }

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

@@ -268,6 +268,10 @@ public class RunManager {
       cmd.add("-Dray.logging.file.path=" + logFile);
     }
 
+    // socket names
+    cmd.add("-Dray.raylet.socket-name=" + rayConfig.rayletSocketName);
+    cmd.add("-Dray.object-store.socket-name=" + rayConfig.objectStoreSocketName);
+
     // Config overwrite
     cmd.add("-Dray.redis.address=" + rayConfig.getRedisAddress());