Configurable Native Library Extraction in Java

java/src/main/java/ai/onnxruntime/OnnxRuntime.java

@@ -48,6 +48,18 @@ final class OnnxRuntime {
   // The initial release of the ORT training API.
   private static final int ORT_TRAINING_API_VERSION_1 = 1;
 
+  /**
+   * The name of the system property which when set gives the path on disk where the ONNX Runtime
+   * native libraries will be extracted to if needed.
+   */
+  static final String ONNXRUNTIME_NATIVE_EXTRACT_PATH = "onnxruntime.native.extract.path";
+
+  /**
+   * The name of the system property which governs whether to delete any extracted libraries on
+   * termination.
+   */
+  static final String ONNXRUNTIME_NATIVE_EXTRACT_CLEANUP = "onnxruntime.native.extract.cleanup";
+
   /**
    * The name of the system property which when set gives the path on disk where the ONNX Runtime
    * native libraries are stored.
@@ -96,8 +108,11 @@ final class OnnxRuntime {
   /** Have the core ONNX Runtime native libraries been loaded */
   private static boolean loaded = false;
 
-  /** The temp directory where native libraries are extracted */
-  private static Path tempDirectory;
+  /** The directory where native libraries are extracted */
+  private static Path extractDirectory;
+
+  /** Whether to delete extracted libraries on termination */
+  private static boolean cleanUp = true;
 
   /** The value of the {@link #ONNXRUNTIME_NATIVE_PATH} system property */
   private static String libraryDirPathProperty;
@@ -171,7 +186,24 @@ static synchronized void init() throws IOException {
     if (loaded) {
       return;
     }
-    tempDirectory = isAndroid() ? null : Files.createTempDirectory("onnxruntime-java");
+
+    cleanUp =
+        Boolean.TRUE
+            .toString()
+            .equalsIgnoreCase(System.getProperty(ONNXRUNTIME_NATIVE_EXTRACT_CLEANUP, "true"));
+
+    if (isAndroid()) {
+      extractDirectory = null;
+    } else {
+      String extractDirectoryProperty = System.getProperty(ONNXRUNTIME_NATIVE_EXTRACT_PATH);
+      if (extractDirectoryProperty != null) {
+        // TODO: Switch this to Path.of when the minimum Java version is 11.
+        extractDirectory = Paths.get(extractDirectoryProperty);
+        Files.createDirectories(extractDirectory);
+      } else {
+        extractDirectory = Files.createTempDirectory("onnxruntime-java");
+      }
+    }
     try {
       libraryDirPathProperty = System.getProperty(ONNXRUNTIME_NATIVE_PATH);
       // Extract and prepare the shared provider library but don't try to load it,
@@ -201,8 +233,8 @@ static synchronized void init() throws IOException {
       version = initialiseVersion();
       loaded = true;
     } finally {
-      if (tempDirectory != null) {
-        cleanUp(tempDirectory.toFile());
+      if (extractDirectory != null && cleanUp) {
+        cleanUp(extractDirectory.toFile());
       }
     }
   }
@@ -429,7 +461,7 @@ private static void load(String library) throws IOException {
   private static Optional<File> extractFromResources(String library) {
     String libraryFileName = mapLibraryName(library);
     String resourcePath = "/ai/onnxruntime/native/" + OS_ARCH_STR + '/' + libraryFileName;
-    File tempFile = tempDirectory.resolve(libraryFileName).toFile();
+    File targetFile = extractDirectory.resolve(libraryFileName).toFile();
     try (InputStream is = OnnxRuntime.class.getResourceAsStream(resourcePath)) {
       if (is == null) {
         // Not found in classpath resources
@@ -443,23 +475,25 @@ private static Optional<File> extractFromResources(String library) {
                 + "' from resource path "
                 + resourcePath
                 + " copying to "
-                + tempFile);
+                + targetFile);
         byte[] buffer = new byte[4096];
         int readBytes;
-        try (FileOutputStream os = new FileOutputStream(tempFile)) {
+        try (FileOutputStream os = new FileOutputStream(targetFile)) {
           while ((readBytes = is.read(buffer)) != -1) {
             os.write(buffer, 0, readBytes);
           }
         }
         logger.log(Level.FINE, "Extracted native library '" + library + "' from resource path");
-        return Optional.of(tempFile);
+        return Optional.of(targetFile);
       }
     } catch (IOException e) {
       logger.log(
           Level.WARNING, "Failed to extract library '" + library + "' from the resources", e);
       return Optional.empty();
     } finally {
-      cleanUp(tempFile);
+      if (cleanUp) {
+        cleanUp(targetFile);
+      }
     }
   }
 

java/src/main/java/ai/onnxruntime/package-info.java

@@ -10,38 +10,43 @@
  * Java (such as fp16) are converted into the nearest Java primitive type when accessed through this
  * API.
  *
- * <p>There are two shared libraries required: <code>onnxruntime</code> and <code>onnxruntime4j_jni
- * </code>. The loader is in {@link ai.onnxruntime.OnnxRuntime} and the logic is in this order:
+ * <p>There are several shared libraries required for operation, primarily <code>onnxruntime</code>
+ * and <code>onnxruntime4j_jni</code>, as well as <code>onnxruntime_providers_shared</code>. The
+ * loader is in {@link ai.onnxruntime.OnnxRuntime} and the logic follows this order:
  *
  * <ol>
  *   <li>The user may signal to skip loading of a shared library using a property in the form <code>
- *       onnxruntime.native.LIB_NAME.skip</code> with a value of <code>true</code>. This means the
- *       user has decided to load the library by some other means.
- *   <li>The user may specify an explicit location of all native library files using a property in
- *       the form <code>onnxruntime.native.path</code>. This uses {@link java.lang.System#load}.
- *   <li>The user may specify an explicit location of the shared library file using a property in
- *       the form <code>onnxruntime.native.LIB_NAME.path</code>. This uses {@link
+ * onnxruntime.native.LIB_NAME.skip</code> with a value of <code>true</code>. This means the user
+ *       has decided to load the library by some other means.
+ *   <li>The user may specify an explicit directory containing all native library files using the
+ *       property <code>onnxruntime.native.path</code>. This uses {@link java.lang.System#load}.
+ *   <li>The user may specify an explicit location of a specific shared library file using a
+ *       property in the form <code>onnxruntime.native.LIB_NAME.path</code>. This uses {@link
  *       java.lang.System#load}.
  *   <li>The shared library is autodiscovered:
  *       <ol>
- *         <li>If the shared library is present in the classpath resources, load using {@link
- *             java.lang.System#load} via a temporary file. Ideally, this should be the default use
- *             case when adding JAR's/dependencies containing the shared libraries to your
- *             classpath.
+ *         <li>If the shared library is present in the classpath resources, it is extracted and
+ *             loaded using {@link java.lang.System#load}.
+ *             <ul>
+ *               <li>The extraction location can be controlled via <code>
+ * onnxruntime.native.extract.path</code>. If unset, a temporary directory is created.
+ *               <li>By default, extracted libraries are deleted on JVM termination. This can be
+ *                   disabled by setting <code>onnxruntime.native.extract.cleanup</code> to <code>
+ * false</code>.
+ *             </ul>
  *         <li>If the shared library is not present in the classpath resources, then load using
- *             {@link java.lang.System#loadLibrary}, which usually looks elsewhere on the filesystem
- *             for the library. The semantics and behavior of that method are system/JVM dependent.
- *             Typically, the <code>java.library.path</code> property is used to specify the
- *             location of native libraries.
+ *             {@link java.lang.System#loadLibrary}, which looks on the standard library paths
+ *             (e.g., <code>java.library.path</code>).
  *       </ol>
  * </ol>
  *
  * For troubleshooting, all shared library loading events are reported to Java logging at the level
  * FINE.
  *
- * <p>Note that CUDA, ROCM, DNNL, OpenVINO and TensorRT are all "shared library execution providers"
- * and must be stored either in the directory containing the ONNX Runtime core native library, or as
- * a classpath resource. This is because these providers are loaded by the ONNX Runtime native
- * library itself and the Java API cannot control the loading location.
+ * <p>Note that CUDA, ROCM, DNNL, OpenVINO, TensorRT, and QNN are all "shared library execution
+ * providers." These, along with WebGPU dependencies like Dawn or DXC, must be stored either in the
+ * directory containing the ONNX Runtime core native library or as a classpath resource. This is
+ * because these providers are loaded by the ONNX Runtime native library itself, and the Java API
+ * handles their extraction/preparation to ensure they are available to the native loader.
  */
 package ai.onnxruntime;