Support reading system configs from native in Java. (ray-project#17703)

* Support reading system configs from native in Java. * Fix lint * Lint cpp * Fix Java cases. * Address comments. * Address comments.
windycom · Aug 12, 2021 · 6d6a1ea · 6d6a1ea
1 parent 623db7c
commit 6d6a1ea
Show file tree

Hide file tree

Showing 7 changed files with 149 additions and 6 deletions.
diff --git a/java/runtime/src/main/java/io/ray/runtime/RayDevRuntime.java b/java/runtime/src/main/java/io/ray/runtime/RayDevRuntime.java
@@ -12,6 +12,11 @@
 import io.ray.runtime.object.LocalModeObjectStore;
 import io.ray.runtime.task.LocalModeTaskExecutor;
 import io.ray.runtime.task.LocalModeTaskSubmitter;
+import io.ray.runtime.util.BinaryFileUtil;
+import io.ray.runtime.util.JniUtils;
+import io.ray.runtime.util.SystemUtil;
+import java.text.SimpleDateFormat;
+import java.util.Date;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -33,6 +38,10 @@ public void start() {
     if (rayConfig.getJobId().isNil()) {
       rayConfig.setJobId(nextJobId());
     }
+
+    updateSessionDir(rayConfig);
+    JniUtils.loadLibrary(rayConfig.sessionDir, BinaryFileUtil.CORE_WORKER_JAVA_LIBRARY, true);
+
     taskExecutor = new LocalModeTaskExecutor(this);
     workerContext = new LocalModeWorkerContext(rayConfig.getJobId());
     objectStore = new LocalModeObjectStore(workerContext);
@@ -118,4 +127,12 @@ private AsyncContext(TaskSpec task) {
       this.task = task;
     }
   }
+
+  private static void updateSessionDir(RayConfig rayConfig) {
+    SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd_hh-mm-ss-ms");
+    Date date = new Date();
+    String sessionDir =
+        String.format("/tmp/ray/session_local_mode_%s_%d", format.format(date), SystemUtil.pid());
+    rayConfig.setSessionDir(sessionDir);
+  }
 }
diff --git a/java/runtime/src/main/java/io/ray/runtime/task/ArgumentsBuilder.java b/java/runtime/src/main/java/io/ray/runtime/task/ArgumentsBuilder.java
@@ -11,6 +11,7 @@
 import io.ray.runtime.object.NativeRayObject;
 import io.ray.runtime.object.ObjectRefImpl;
 import io.ray.runtime.object.ObjectSerializer;
+import io.ray.runtime.util.SystemConfig;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -23,8 +24,8 @@ public class ArgumentsBuilder {
    * If the the size of an argument's serialized data is smaller than this number, the argument will
    * be passed by value. Otherwise it'll be passed by reference.
    */
-  // TODO(kfstorm): Read from internal config `max_direct_call_object_size`.
-  public static final int LARGEST_SIZE_PASS_BY_VALUE = 100 * 1024;
+  public static final long LARGEST_SIZE_PASS_BY_VALUE =
+      ((Double) SystemConfig.get("max_direct_call_object_size")).longValue();
 
   /** This dummy type is also defined in signature.py. Please keep it synced. */
   private static final NativeRayObject PYTHON_DUMMY_TYPE =

diff --git a/java/runtime/src/main/java/io/ray/runtime/util/SystemConfig.java b/java/runtime/src/main/java/io/ray/runtime/util/SystemConfig.java
@@ -0,0 +1,27 @@
+package io.ray.runtime.util;
+
+import com.google.common.base.Preconditions;
+import com.google.gson.Gson;
+import java.util.HashMap;
+
+/** The utility class to read system config from native code. */
+public class SystemConfig {
+
+  private static Gson gson = new Gson();
+
+  /// A cache to avoid duplicated reading via JNI.
+  private static HashMap<String, Object> cachedConfigs = new HashMap<>();
+
+  public static synchronized Object get(String key) {
+    if (cachedConfigs.containsKey(key)) {
+      return cachedConfigs.get(key);
+    }
+
+    Object val = gson.fromJson(nativeGetSystemConfig(key), Object.class);
+    Preconditions.checkNotNull(val);
+    cachedConfigs.put(key, val);
+    return val;
+  }
+
+  private static native String nativeGetSystemConfig(String key);
+}
diff --git a/java/test/src/main/java/io/ray/test/RayCallTest.java b/java/test/src/main/java/io/ray/test/RayCallTest.java
@@ -16,8 +16,14 @@
 /** Test Ray.call API */
 public class RayCallTest extends BaseTest {
 
-  private static final byte[] LARGE_RAW_DATA =
-      new byte[ArgumentsBuilder.LARGEST_SIZE_PASS_BY_VALUE + 100];
+  private static byte[] LARGE_RAW_DATA = null;
+
+  private static byte[] getLargeRawData() {
+    if (LARGE_RAW_DATA == null) {
+      LARGE_RAW_DATA = new byte[(int) ArgumentsBuilder.LARGEST_SIZE_PASS_BY_VALUE + 100];
+    }
+    return LARGE_RAW_DATA;
+  }
 
   private static int testInt(int val) {
     return val;
@@ -155,11 +161,11 @@ public void testNumberOfParameters() {
   }
 
   private static Boolean testLargeRawData(byte[] data) {
-    return Arrays.equals(data, LARGE_RAW_DATA);
+    return Arrays.equals(data, getLargeRawData());
   }
 
   @Test
   public void testLargeRawDataArgument() {
-    Assert.assertTrue(Ray.task(RayCallTest::testLargeRawData, LARGE_RAW_DATA).remote().get());
+    Assert.assertTrue(Ray.task(RayCallTest::testLargeRawData, getLargeRawData()).remote().get());
   }
 }
diff --git a/java/test/src/main/java/io/ray/test/SystemConfigTest.java b/java/test/src/main/java/io/ray/test/SystemConfigTest.java
@@ -0,0 +1,14 @@
+package io.ray.test;
+
+import io.ray.runtime.util.SystemConfig;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+@Test
+public class SystemConfigTest extends BaseTest {
+
+  public void testDefaultConfigs() {
+    long ret = ((Double) SystemConfig.get("max_direct_call_object_size")).longValue();
+    Assert.assertEquals(ret, 100 * 1024);
+  }
+}
diff --git a/src/ray/core_worker/lib/java/io_ray_runtime_util_SystemConfig.cc b/src/ray/core_worker/lib/java/io_ray_runtime_util_SystemConfig.cc
@@ -0,0 +1,41 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "io_ray_runtime_util_SystemConfig.h"
+
+#include <jni.h>
+
+#include "jni_utils.h"
+#include "nlohmann/json.hpp"
+#include "ray/util/logging.h"
+
+using json = nlohmann::json;
+
+JNIEXPORT jstring JNICALL Java_io_ray_runtime_util_SystemConfig_nativeGetSystemConfig(
+    JNIEnv *env, jclass clz, jstring java_key) {
+  RAY_CHECK(java_key != nullptr);
+  const auto key = JavaStringToNativeString(env, java_key);
+
+  /// -----------Include ray_config_def.h to set config items.-------------------
+#define RAY_CONFIG(type, name, default_value)   \
+  if (key == #name) {                           \
+    json j = RayConfig::instance().name();      \
+    return env->NewStringUTF(j.dump().c_str()); \
+  }
+
+#include "ray/common/ray_config_def.h"
+/// ---------------------------------------------------------------------
+#undef RAY_CONFIG
+  RAY_LOG(FATAL) << "Unsupported system config: " << key;
+  return nullptr;
+}
diff --git a/src/ray/core_worker/lib/java/io_ray_runtime_util_SystemConfig.h b/src/ray/core_worker/lib/java/io_ray_runtime_util_SystemConfig.h
@@ -0,0 +1,37 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/* DO NOT EDIT THIS FILE - it is machine generated */
+
+#include <jni.h>
+
+/* Header for class io_ray_runtime_util_SystemConfig */
+#ifndef _Included_io_ray_runtime_util_SystemConfig
+#define _Included_io_ray_runtime_util_SystemConfig
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*
+ * Class:     io_ray_runtime_util_SystemConfig
+ * Method:    nativeGetSystemConfig
+ * Signature: (Ljava/lang/String;)Ljava/lang/String;
+ */
+JNIEXPORT jstring JNICALL
+Java_io_ray_runtime_util_SystemConfig_nativeGetSystemConfig(JNIEnv *, jclass, jstring);
+#ifdef __cplusplus
+}
+#endif
+#endif