[java] Support resources management in raylet mode. (#2606)

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

@@ -16,6 +16,7 @@ import org.ray.core.model.RunMode;
 import org.ray.runner.RunInfo.ProcessType;
 import org.ray.spi.PathConfig;
 import org.ray.spi.model.AddressInfo;
+import org.ray.util.ResourceUtil;
 import org.ray.util.StringUtil;
 import org.ray.util.config.ConfigReader;
 import org.ray.util.logger.RayLog;
@@ -350,10 +351,13 @@ public class RunManager {
       startObjectStore(0, info, params.working_directory + "/store",
           params.redis_address, params.node_ip_address, params.redirect, params.cleanup);
 
+      Map<String, Double> staticResources =
+          ResourceUtil.getResourcesMapFromString(params.static_resources);
+
       //Start raylet
-      startRaylet(storeName, info, params.num_cpus[0],params.num_gpus[0],
-          params.num_workers,params.working_directory + "/raylet",
-          params.redis_address, params.node_ip_address, params.redirect, params.cleanup);
+      startRaylet(storeName, info, params.num_workers,
+          params.working_directory + "/raylet", params.redis_address,
+          params.node_ip_address, params.redirect, staticResources, params.cleanup);
 
       runInfo.localStores.add(info);
     } else {
@@ -677,10 +681,9 @@ public class RunManager {
     }
   }
 
-  private void startRaylet(String storeName, AddressInfo info, int numCpus,
-                           int numGpus, int numWorkers, String workDir,
-                           String redisAddress, String ip, boolean redirect,
-                           boolean cleanup) {
+  private void startRaylet(String storeName, AddressInfo info, int numWorkers,
+                           String workDir, String redisAddress, String ip, boolean redirect,
+                           Map<String, Double> staticResources, boolean cleanup) {
 
     int rpcPort = params.raylet_port;
     String rayletSocketName = "/tmp/raylet" + rpcPort;
@@ -695,8 +698,8 @@ public class RunManager {
     assert (sep != -1);
     String gcsIp = redisAddress.substring(0, sep);
     String gcsPort = redisAddress.substring(sep + 1);
-    
-    String resourceArgument = "GPU," + numGpus + ",CPU," + numCpus;
+
+    String resourceArgument = ResourceUtil.getResourcesStringFromMap(staticResources);
 
     String[] cmds = new String[]{filePath, rayletSocketName, storeName, ip, gcsIp,
                                  gcsPort, "" + numWorkers, workerCmd, resourceArgument};

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

@@ -5,11 +5,13 @@ import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import org.ray.api.UniqueID;
 import org.ray.core.RayRuntime;
 import org.ray.spi.LocalSchedulerLink;
 import org.ray.spi.model.FunctionArg;
 import org.ray.spi.model.TaskSpec;
+import org.ray.util.ResourceUtil;
 import org.ray.util.logger.RayLog;
 
 /**
@@ -64,6 +66,20 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
 
   @Override
   public void submitTask(TaskSpec task) {
+    // We don't support resources management in non raylet mode.
+    if (!useRaylet) {
+      task.resources.clear();
+      task.resources.put(ResourceUtil.CPU_LITERAL, 0.0);
+    } else {
+      if (!task.resources.containsKey(ResourceUtil.CPU_LITERAL)) {
+        task.resources.put(ResourceUtil.CPU_LITERAL, 0.0);
+      }
+
+      if (!task.resources.containsKey(ResourceUtil.GPU_LITERAL)) {
+        task.resources.put(ResourceUtil.GPU_LITERAL, 0.0);
+      }
+    }
+
     ByteBuffer info = taskSpec2Info(task);
     byte[] a = null;
     if (!task.actorId.isNil()) {
@@ -220,13 +236,15 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     // The required_resources vector indicates the quantities of the different
     // resources required by this task. The index in this vector corresponds to
     // the resource type defined in the ResourceIndex enum. For example,
-
-    int[]requiredResourcesOffsets = new int[1];
-    for (int i = 0; i < requiredResourcesOffsets.length; i++) {
-      int keyOffset = 0;
-      keyOffset = fbb.createString(ByteBuffer.wrap("CPU".getBytes()));
-      requiredResourcesOffsets[i] = ResourcePair.createResourcePair(fbb, keyOffset, 0.0);
+    int[] requiredResourcesOffsets = new int[task.resources.size()];
+    int i = 0;
+    for (Map.Entry<String, Double> entry : task.resources.entrySet()) {
+      int keyOffset = fbb.createString(ByteBuffer.wrap(entry.getKey().getBytes()));
+      requiredResourcesOffsets[i] =
+          ResourcePair.createResourcePair(fbb, keyOffset, entry.getValue());
+      i++;
     }
+
     int requiredResourcesOffset = fbb.createVectorOfTables(requiredResourcesOffsets);
 
     int root = TaskInfo.createTaskInfo(