[Java] Single-process mode (ray-project#4245)

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

@@ -66,7 +66,8 @@ public AbstractRayRuntime(RayConfig rayConfig) {
     this.rayConfig = rayConfig;
     functionManager = new FunctionManager(rayConfig.driverResourcePath);
     worker = new Worker(this);
-    workerContext = new WorkerContext(rayConfig.workerMode, rayConfig.driverId);
+    workerContext = new WorkerContext(rayConfig.workerMode,
+            rayConfig.driverId, rayConfig.runMode);
   }
 
   /**

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

@@ -17,12 +17,12 @@ public RayDevRuntime(RayConfig rayConfig) {
   public void start() {
     store = new MockObjectStore(this);
     objectStoreProxy = new ObjectStoreProxy(this, null);
-    rayletClient = new MockRayletClient(this, store);
+    rayletClient = new MockRayletClient(this, rayConfig.numberExecThreadsForDevRuntime);
   }
 
   @Override
   public void shutdown() {
-    // nothing to do
+    rayletClient.destroy();
   }
 
   public MockObjectStore getObjectStore() {

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

@@ -2,6 +2,7 @@
 
 import com.google.common.base.Preconditions;
 import org.ray.api.id.UniqueId;
+import org.ray.runtime.config.RunMode;
 import org.ray.runtime.config.WorkerMode;
 import org.ray.runtime.task.TaskSpec;
 import org.slf4j.Logger;
@@ -34,12 +35,17 @@ public class WorkerContext {
    */
   private long mainThreadId;
 
+  /**
+   * The run-mode of this worker.
+   */
+  private RunMode runMode;
 
-  public WorkerContext(WorkerMode workerMode, UniqueId driverId) {
+  public WorkerContext(WorkerMode workerMode, UniqueId driverId, RunMode runMode) {
     mainThreadId = Thread.currentThread().getId();
     taskIndex = ThreadLocal.withInitial(() -> 0);
     putIndex = ThreadLocal.withInitial(() -> 0);
     currentTaskId = ThreadLocal.withInitial(UniqueId::randomId);
+    this.runMode = runMode;
     currentClassLoader = null;
     if (workerMode == WorkerMode.DRIVER) {
       workerId = driverId;
@@ -65,10 +71,12 @@ public UniqueId getCurrentTaskId() {
    * be called from the main thread.
    */
   public void setCurrentTask(TaskSpec task, ClassLoader classLoader) {
-    Preconditions.checkState(
-        Thread.currentThread().getId() == mainThreadId,
-        "This method should only be called from the main thread."
-    );
+    if (runMode == RunMode.CLUSTER) {
+      Preconditions.checkState(
+              Thread.currentThread().getId() == mainThreadId,
+              "This method should only be called from the main thread."
+      );
+    }
 
     Preconditions.checkNotNull(task);
     this.currentTaskId.set(task.taskId);

java/runtime/src/main/java/org/ray/runtime/config/RayConfig.java

@@ -63,6 +63,11 @@ public class RayConfig {
   public final String driverResourcePath;
   public final String pythonWorkerCommand;
 
+  /**
+   * Number of threads that execute tasks.
+   */
+  public final int numberExecThreadsForDevRuntime;
+
   private void validate() {
     if (workerMode == WorkerMode.WORKER) {
       Preconditions.checkArgument(redisAddress != null,
@@ -196,6 +201,9 @@ public RayConfig(Config config) {
       driverResourcePath = null;
     }
 
+    // Number of threads that execute tasks.
+    numberExecThreadsForDevRuntime = config.getInt("ray.dev-runtime.execution-parallelism");
+
     // validate config
     validate();
     LOGGER.debug("Created config: {}", this);

java/runtime/src/main/java/org/ray/runtime/objectstore/MockObjectStore.java

@@ -6,11 +6,12 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+
 import org.apache.arrow.plasma.ObjectStoreLink;
-import org.apache.arrow.plasma.ObjectStoreLink.ObjectStoreData;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.RayDevRuntime;
-import org.ray.runtime.raylet.MockRayletClient;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -20,13 +21,21 @@
 public class MockObjectStore implements ObjectStoreLink {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(MockObjectStore.class);
+
+  private static final int GET_CHECK_INTERVAL_MS = 100;
+
   private final RayDevRuntime runtime;
   private final Map<UniqueId, byte[]> data = new ConcurrentHashMap<>();
   private final Map<UniqueId, byte[]> metadata = new ConcurrentHashMap<>();
-  private MockRayletClient scheduler = null;
+  private final List<Consumer<UniqueId>> objectPutCallbacks;
 
   public MockObjectStore(RayDevRuntime runtime) {
     this.runtime = runtime;
+    this.objectPutCallbacks = new ArrayList<>();
+  }
+
+  public void addObjectPutCallback(Consumer<UniqueId> callback) {
+    this.objectPutCallbacks.add(callback);
   }
 
   @Override
@@ -41,34 +50,56 @@ public void put(byte[] objectId, byte[] value, byte[] metadataValue) {
     if (metadataValue != null) {
       metadata.put(uniqueId, metadataValue);
     }
-    if (scheduler != null) {
-      scheduler.onObjectPut(uniqueId);
+    UniqueId id = new UniqueId(objectId);
+    for (Consumer<UniqueId> callback : objectPutCallbacks) {
+      callback.accept(id);
     }
   }
 
+  @Override
+  public byte[] get(byte[] objectId, int timeoutMs, boolean isMetadata) {
+    return get(new byte[][] {objectId}, timeoutMs, isMetadata).get(0);
+  }
+
   @Override
   public List<byte[]> get(byte[][] objectIds, int timeoutMs, boolean isMetadata) {
-    final Map<UniqueId, byte[]> dataMap = isMetadata ? metadata : data;
-    ArrayList<byte[]> rets = new ArrayList<>(objectIds.length);
-    for (byte[] objId : objectIds) {
-      UniqueId uniqueId = new UniqueId(objId);
-      LOGGER.info("{} is notified for objectid {}",logPrefix(), uniqueId);
-      rets.add(dataMap.get(uniqueId));
-    }
-    return rets;
+    return get(objectIds, timeoutMs)
+            .stream()
+            .map(data -> isMetadata ? data.data : data.metadata)
+            .collect(Collectors.toList());
   }
 
   @Override
   public List<ObjectStoreData> get(byte[][] objectIds, int timeoutMs) {
+    int ready = 0;
+    int remainingTime = timeoutMs;
+    boolean firstCheck = true;
+    while (ready < objectIds.length && remainingTime > 0) {
+      if (!firstCheck) {
+        int sleepTime = Math.min(remainingTime, GET_CHECK_INTERVAL_MS);
+        try {
+          Thread.sleep(sleepTime);
+        } catch (InterruptedException e) {
+          LOGGER.warn("Got InterruptedException while sleeping.");
+        }
+        remainingTime -= sleepTime;
+      }
+      ready = 0;
+      for (byte[] id : objectIds) {
+        if (data.containsKey(new UniqueId(id))) {
+          ready += 1;
+        }
+      }
+      firstCheck = false;
+    }
     ArrayList<ObjectStoreData> rets = new ArrayList<>();
-    // TODO(yuhguo): make ObjectStoreData's constructor public.
-    for (byte[] objId : objectIds) {
-      UniqueId uniqueId = new UniqueId(objId);
+    for (byte[] id : objectIds) {
       try {
         Constructor<ObjectStoreData> constructor = ObjectStoreData.class.getConstructor(
-            byte[].class, byte[].class);
+                byte[].class, byte[].class);
         constructor.setAccessible(true);
-        rets.add(constructor.newInstance(metadata.get(uniqueId), data.get(uniqueId)));
+        rets.add(constructor.newInstance(metadata.get(new UniqueId(id)),
+                data.get(new UniqueId(id))));
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
@@ -119,7 +150,8 @@ public boolean isObjectReady(UniqueId id) {
     return data.containsKey(id);
   }
 
-  public void registerScheduler(MockRayletClient s) {
-    scheduler = s;
+  public void free(UniqueId id) {
+    data.remove(id);
+    metadata.remove(id);
   }
 }