[Java] Allow passing internal config from raylet to Java worker (#7532)

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

@@ -61,7 +61,7 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
     } catch (IOException e) {
       throw new RuntimeException("Failed to create the log directory.", e);
     }
-    nativeSetup(rayConfig.logDir);
+    nativeSetup(rayConfig.logDir, rayConfig.rayletConfigParameters);
     Runtime.getRuntime().addShutdownHook(new Thread(RayNativeRuntime::nativeShutdownHook));
   }
 
@@ -193,7 +193,7 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
 
   private static native void nativeDestroyCoreWorker(long nativeCoreWorkerPointer);
 
-  private static native void nativeSetup(String logDir);
+  private static native void nativeSetup(String logDir, Map<String, String> rayletConfigParameters);
 
   private static native void nativeShutdownHook();
 

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

@@ -10,7 +10,7 @@ import com.typesafe.config.ConfigValue;
 import java.io.File;
 import java.time.LocalDateTime;
 import java.time.format.DateTimeFormatter;
-import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
@@ -67,7 +67,7 @@ public class RayConfig {
 
   public String rayletSocketName;
   private int nodeManagerPort;
-  public final List<String> rayletConfigParameters;
+  public final Map<String, String> rayletConfigParameters;
 
   public final String jobResourcePath;
   public final String pythonWorkerCommand;
@@ -204,11 +204,11 @@ public class RayConfig {
     }
 
     // Raylet parameters.
-    rayletConfigParameters = new ArrayList<>();
+    rayletConfigParameters = new HashMap<>();
     Config rayletConfig = config.getConfig("ray.raylet.config");
     for (Map.Entry<String, ConfigValue> entry : rayletConfig.entrySet()) {
-      String parameter = entry.getKey() + "," + entry.getValue().unwrapped();
-      rayletConfigParameters.add(parameter);
+      Object value = entry.getValue().unwrapped();
+      rayletConfigParameters.put(entry.getKey(), value == null ? "" : value.toString());
     }
 
     // Job resource path.

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

@@ -240,7 +240,10 @@ public class RunManager {
           gcsServerFile.getAbsolutePath(),
           String.format("--redis_address=%s", rayConfig.getRedisIp()),
           String.format("--redis_port=%d", rayConfig.getRedisPort()),
-          String.format("--config_list=%s", String.join(",", rayConfig.rayletConfigParameters)),
+          String.format("--config_list=%s",
+              rayConfig.rayletConfigParameters.entrySet().stream()
+                  .map(entry -> entry.getKey() + "," + entry.getValue()).collect(Collectors
+                      .joining(","))),
           String.format("--redis_password=%s", redisPasswordOption)
       );
       startProcess(command, null, "gcs_server");
@@ -316,7 +319,9 @@ public class RunManager {
         String.format("--maximum_startup_concurrency=%d", maximumStartupConcurrency),
         String.format("--static_resource_list=%s",
             ResourceUtil.getResourcesStringFromMap(rayConfig.resources)),
-        String.format("--config_list=%s", String.join(",", rayConfig.rayletConfigParameters)),
+        String.format("--config_list=%s", rayConfig.rayletConfigParameters.entrySet().stream()
+            .map(entry -> entry.getKey() + "," + entry.getValue())
+            .collect(Collectors.joining(","))),
         String.format("--python_worker_command=%s", buildPythonWorkerCommand()),
         String.format("--java_worker_command=%s", buildWorkerCommand()),
         String.format("--redis_password=%s", redisPasswordOption)
@@ -378,8 +383,8 @@ public class RunManager {
       cmd.add("-Dray.redis.password=" + rayConfig.headRedisPassword);
     }
 
-    // Number of workers per Java worker process
-    cmd.add("-Dray.raylet.config.num_workers_per_process_java=RAY_WORKER_NUM_WORKERS_PLACEHOLDER");
+
+    cmd.add("RAY_WORKER_RAYLET_CONFIG_PLACEHOLDER");
 
     cmd.addAll(rayConfig.jvmParameters);
 

java/test/src/main/java/org/ray/api/test/BaseMultiLanguageTest.java

@@ -1,6 +1,5 @@
 package org.ray.api.test;
 
-import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.gson.Gson;
@@ -13,6 +12,7 @@ import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.ray.api.Ray;
+import org.ray.runtime.config.RayConfig;
 import org.ray.runtime.util.NetworkUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -90,15 +90,9 @@ public abstract class BaseMultiLanguageTest {
         String.format("--node-manager-port=%s", nodeManagerPort),
         "--load-code-from-local",
         "--include-java",
-        "--java-worker-options=" + workerOptions
+        "--java-worker-options=" + workerOptions,
+        "--internal-config=" + new Gson().toJson(RayConfig.create().rayletConfigParameters)
     );
-    String numWorkersPerProcessJava = System
-        .getProperty("ray.raylet.config.num_workers_per_process_java");
-    if (!Strings.isNullOrEmpty(numWorkersPerProcessJava)) {
-      startCommand = ImmutableList.<String>builder().addAll(startCommand)
-          .add(String.format("--internal-config={\"num_workers_per_process_java\": %s}",
-              numWorkersPerProcessJava)).build();
-    }
     if (!executeCommand(startCommand, 10, getRayStartEnv())) {
       throw new RuntimeException("Couldn't start ray cluster.");
     }

java/test/src/main/java/org/ray/api/test/RayletConfigTest.java

@@ -0,0 +1,39 @@
+package org.ray.api.test;
+
+import org.ray.api.Ray;
+import org.ray.api.RayActor;
+import org.ray.api.TestUtils;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+public class RayletConfigTest extends BaseTest {
+
+  private static final String RAY_CONFIG_KEY = "num_workers_per_process_java";
+  private static final String RAY_CONFIG_VALUE = "2";
+
+  @BeforeClass
+  public void beforeClass() {
+    System.setProperty("ray.raylet.config." + RAY_CONFIG_KEY, RAY_CONFIG_VALUE);
+  }
+
+  @AfterClass
+  public void afterClass() {
+    System.clearProperty("ray.raylet.config." + RAY_CONFIG_KEY);
+  }
+
+  public static class TestActor {
+
+    public String getConfigValue() {
+      return TestUtils.getRuntime().getRayConfig().rayletConfigParameters.get(RAY_CONFIG_KEY);
+    }
+  }
+
+  @Test
+  public void testRayletConfigPassThrough() {
+    RayActor<TestActor> actor = Ray.createActor(TestActor::new);
+    String configValue = actor.call(TestActor::getConfigValue).get();
+    Assert.assertEquals(configValue, RAY_CONFIG_VALUE);
+  }
+}