[Streaming] Streaming data transfer java (#6474)

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

@@ -1,10 +1,8 @@
 package org.ray.runtime;
 
 import com.google.common.base.Preconditions;
-import com.google.common.base.Strings;
 import java.io.File;
 import java.io.IOException;
-import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.commons.io.FileUtils;
@@ -22,7 +20,7 @@ import org.ray.runtime.runner.RunManager;
 import org.ray.runtime.task.NativeTaskExecutor;
 import org.ray.runtime.task.NativeTaskSubmitter;
 import org.ray.runtime.task.TaskExecutor;
-import org.ray.runtime.util.FileUtil;
+import org.ray.runtime.util.JniUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -42,16 +40,11 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
 
   static {
     LOGGER.debug("Loading native libraries.");
-    // Load native libraries.
-    String[] libraries = new String[]{"core_worker_library_java"};
-    for (String library : libraries) {
-      String fileName = System.mapLibraryName(library);
-      try (FileUtil.TempFile libFile = FileUtil.getTempFileFromResource(fileName)) {
-        System.load(libFile.getFile().getAbsolutePath());
-      }
-      LOGGER.debug("Native libraries loaded.");
-    }
-
+    // Expose ray ABI symbols which may be depended by other shared
+    // libraries such as libstreaming_java.so.
+    // See BUILD.bazel:libcore_worker_library_java.so
+    JniUtils.loadLibrary("core_worker_library_java", true);
+    LOGGER.debug("Native libraries loaded.");
     RayConfig globalRayConfig = RayConfig.create();
     resetLibraryPath(globalRayConfig);
 
@@ -65,30 +58,9 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
   }
 
   private static void resetLibraryPath(RayConfig rayConfig) {
-    if (rayConfig.libraryPath.isEmpty()) {
-      return;
-    }
-
-    String path = System.getProperty("java.library.path");
-    if (Strings.isNullOrEmpty(path)) {
-      path = "";
-    } else {
-      path += ":";
-    }
-    path += String.join(":", rayConfig.libraryPath);
-
-    // This is a hack to reset library path at runtime,
-    // see https://stackoverflow.com/questions/15409223/.
-    System.setProperty("java.library.path", path);
-    // Set sys_paths to null so that java.library.path will be re-evaluated next time it is needed.
-    final Field sysPathsField;
-    try {
-      sysPathsField = ClassLoader.class.getDeclaredField("sys_paths");
-      sysPathsField.setAccessible(true);
-      sysPathsField.set(null, null);
-    } catch (NoSuchFieldException | IllegalAccessException e) {
-      LOGGER.error("Failed to set library path.", e);
-    }
+    String separator = System.getProperty("path.separator");
+    String libraryPath = String.join(separator, rayConfig.libraryPath);
+    JniUtils.resetLibraryPath(libraryPath);
   }
 
   public RayNativeRuntime(RayConfig rayConfig, FunctionManager functionManager) {

java/runtime/src/main/java/org/ray/runtime/runner/worker/DefaultWorker.java

@@ -17,6 +17,9 @@ public class DefaultWorker {
   public static void main(String[] args) {
     try {
       System.setProperty("ray.worker.mode", "WORKER");
+      // Set run-mode to `CLUSTER` explicitly, to prevent the DefaultWorker to receive
+      // a wrong run-mode parameter through jvm options.
+      System.setProperty("ray.run-mode", "CLUSTER");
       Thread.setDefaultUncaughtExceptionHandler((Thread t, Throwable e) -> {
         LOGGER.error("Uncaught worker exception in thread {}: {}", t, e);
       });

java/runtime/src/main/java/org/ray/runtime/util/JniUtils.java

@@ -0,0 +1,84 @@
+package org.ray.runtime.util;
+
+import com.google.common.base.Strings;
+import com.google.common.collect.Sets;
+import com.sun.jna.NativeLibrary;
+import java.lang.reflect.Field;
+import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class JniUtils {
+  private static final Logger LOGGER = LoggerFactory.getLogger(JniUtils.class);
+  private static Set<String> loadedLibs = Sets.newHashSet();
+
+  /**
+   * Loads the native library specified by the <code>libraryName</code> argument.
+   * The <code>libraryName</code> argument must not contain any platform specific
+   * prefix, file extension or path.
+   *
+   * @param libraryName the name of the library.
+   */
+  public static synchronized void loadLibrary(String libraryName) {
+    loadLibrary(libraryName, false);
+  }
+
+  /**
+   * Loads the native library specified by the <code>libraryName</code> argument.
+   * The <code>libraryName</code> argument must not contain any platform specific
+   * prefix, file extension or path.
+   *
+   * @param libraryName   the name of the library.
+   * @param exportSymbols export symbols of library so that it can be used by other libs.
+   */
+  public static synchronized void loadLibrary(String libraryName, boolean exportSymbols) {
+    if (!loadedLibs.contains(libraryName)) {
+      LOGGER.debug("Loading native library {}.", libraryName);
+      // Load native library.
+      String fileName = System.mapLibraryName(libraryName);
+      String libPath = null;
+      try (FileUtil.TempFile libFile = FileUtil.getTempFileFromResource(fileName)) {
+        libPath = libFile.getFile().getAbsolutePath();
+        if (exportSymbols) {
+          // Expose library symbols using RTLD_GLOBAL which may be depended by other shared
+          // libraries.
+          NativeLibrary.getInstance(libFile.getFile().getAbsolutePath());
+        }
+        System.load(libPath);
+      }
+      LOGGER.debug("Native library loaded.");
+      resetLibraryPath(libPath);
+      loadedLibs.add(libraryName);
+    }
+  }
+
+  /**
+   * This is a hack to reset library path at runtime. Please don't use it outside of ray
+   */
+  public static synchronized void resetLibraryPath(String libPath) {
+    if (Strings.isNullOrEmpty(libPath)) {
+      return;
+    }
+    String path = System.getProperty("java.library.path");
+    String separator = System.getProperty("path.separator");
+    if (Strings.isNullOrEmpty(path)) {
+      path = "";
+    } else {
+      path += separator;
+    }
+    path += String.join(separator, libPath);
+
+    // This is a hack to reset library path at runtime,
+    // see https://stackoverflow.com/questions/15409223/.
+    System.setProperty("java.library.path", path);
+    // Set sys_paths to null so that java.library.path will be re-evaluated next time it is needed.
+    final Field sysPathsField;
+    try {
+      sysPathsField = ClassLoader.class.getDeclaredField("sys_paths");
+      sysPathsField.setAccessible(true);
+      sysPathsField.set(null, null);
+    } catch (NoSuchFieldException | IllegalAccessException e) {
+      LOGGER.error("Failed to set library path.", e);
+    }
+  }
+}