Repository: deepjavalibrary/djl
Pull request #2882 (2023-12-02 00:37:37)

[api] Use folk java process to avoid jvm consume GPU memory
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1363 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`606d4f7`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.19%.
> Report is 931 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...va/ai/djl/modality/nlp/generate/TextGenerator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9UZXh0R2VuZXJhdG9yLmphdmE=) | 2.81% | [276 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/nlp/generate/SeqBatcher.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaGVyLmphdmE=) | 0.75% | [132 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...ity/nlp/generate/ContrastiveSeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9Db250cmFzdGl2ZVNlcUJhdGNoU2NoZWR1bGVyLmphdmE=) | 2.97% | [98 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/djl/modality/nlp/generate/SeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaFNjaGVkdWxlci5qYXZh) | 9.83% | [55 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/cv/BufferedImageFactory.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0J1ZmZlcmVkSW1hZ2VGYWN0b3J5LmphdmE=) | 40.96% | [47 Missing and 2 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...a/ai/djl/modality/nlp/generate/StepGeneration.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TdGVwR2VuZXJhdGlvbi5qYXZh) | 2.04% | [48 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [api/src/main/java/ai/djl/ndarray/NDArray.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheS5qYXZh) | 43.42% | [39 Missing and 4 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...n/java/ai/djl/modality/cv/output/CategoryMask.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9DYXRlZ29yeU1hc2suamF2YQ==) | 22.00% | [39 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheUFkYXB0ZXIuamF2YQ==) | 71.21% | [31 Missing and 7 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | 37.50% | [35 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| ... and [76 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2882      +/-   ##
============================================
+ Coverage     72.08%   72.19%   +0.10%     
- Complexity     5126     7173    +2047     
============================================
  Files           473      708     +235     
  Lines         21970    31988   +10018     
  Branches       2351     3332     +981     
============================================
+ Hits          15838    23095    +7257     
- Misses         4925     7305    +2380     
- Partials       1207     1588     +381     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2882?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
index b0b8e3e42..edcbf40ee 100644
--- a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
+++ b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
@@ -22,7 +22,11 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
 import java.lang.management.MemoryUsage;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Locale;
 import java.util.regex.Pattern;
 
@@ -33,6 +37,8 @@ public final class CudaUtils {
 
     private static final CudaLibrary LIB = loadLibrary();
 
+    private static String[] gpuInfo;
+
     private CudaUtils() {}
 
     /**
@@ -49,7 +55,15 @@ public final class CudaUtils {
      *
      * @return the number of GPUs available in the system
      */
+    @SuppressWarnings("PMD.NonThreadSafeSingleton")
     public static int getGpuCount() {
+        if (Boolean.getBoolean("ai.djl.util.cuda.folk")) {
+            if (gpuInfo == null) {
+                gpuInfo = execute(-1); // NOPMD
+            }
+            return Integer.parseInt(gpuInfo[0]);
+        }
+
         if (LIB == null) {
             return 0;
         }
@@ -79,7 +93,19 @@ public final class CudaUtils {
      *
      * @return the version of CUDA runtime
      */
+    @SuppressWarnings("PMD.NonThreadSafeSingleton")
     public static int getCudaVersion() {
+        if (Boolean.getBoolean("ai.djl.util.cuda.folk")) {
+            if (gpuInfo == null) {
+                gpuInfo = execute(-1);
+            }
+            int version = Integer.parseInt(gpuInfo[1]);
+            if (version == -1) {
+                throw new IllegalArgumentException("No cuda device found.");
+            }
+            return version;
+        }
+
         if (LIB == null) {
             throw new IllegalStateException("No cuda library is loaded.");
         }
@@ -95,9 +121,6 @@ public final class CudaUtils {
      * @return the version string of CUDA runtime
      */
     public static String getCudaVersionString() {
-        if (LIB == null) {
-            throw new IllegalStateException("No cuda library is loaded.");
-        }
         int version = getCudaVersion();
         int major = version / 1000;
         int minor = (version / 10) % 10;
@@ -111,6 +134,14 @@ public final class CudaUtils {
      * @return the CUDA compute capability
      */
     public static String getComputeCapability(int device) {
+        if (Boolean.getBoolean("ai.djl.util.cuda.folk")) {
+            String[] ret = execute(device);
+            if (ret.length != 3) {
+                throw new IllegalArgumentException(ret[0]);
+            }
+            return ret[0];
+        }
+
         if (LIB == null) {
             throw new IllegalStateException("No cuda library is loaded.");
         }
@@ -137,6 +168,16 @@ public final class CudaUtils {
             throw new IllegalArgumentException("Only GPU device is allowed.");
         }
 
+        if (Boolean.getBoolean("ai.djl.util.cuda.folk")) {
+            String[] ret = execute(device.getDeviceId());
+            if (ret.length != 3) {
+                throw new IllegalArgumentException(ret[0]);
+            }
+            long total = Long.parseLong(ret[1]);
+            long used = Long.parseLong(ret[2]);
+            return new MemoryUsage(-1, used, used, total);
+        }
+
         if (LIB == null) {
             throw new IllegalStateException("No GPU device detected.");
         }
@@ -155,8 +196,42 @@ public final class CudaUtils {
         return new MemoryUsage(-1, committed, committed, total[0]);
     }
 
+    /**
+     * The main entrypoint to get CUDA information with command line.
+     *
+     * @param args the command line arguments.
+     */
+    @SuppressWarnings("PMD.SystemPrintln")
+    public static void main(String[] args) {
+        int gpuCount = getGpuCount();
+        if (args.length == 0) {
+            if (gpuCount <= 0) {
+                System.out.println("0,-1");
+                return;
+            }
+            int cudaVersion = getCudaVersion();
+            System.out.println(gpuCount + "," + cudaVersion);
+            return;
+        }
+        try {
+            int deviceId = Integer.parseInt(args[0]);
+            if (deviceId < 0 || deviceId >= gpuCount) {
+                System.out.println("Invalid device: " + deviceId);
+                return;
+            }
+            MemoryUsage mem = getGpuMemory(Device.gpu(deviceId));
+            String cc = getComputeCapability(deviceId);
+            System.out.println(cc + ',' + mem.getMax() + ',' + mem.getUsed());
+        } catch (NumberFormatException e) {
+            System.out.println("Invalid device: " + args[0]);
+        }
+    }
+
     private static CudaLibrary loadLibrary() {
         try {
+            if (Boolean.getBoolean("ai.djl.util.cuda.folk")) {
+                return null;
+            }
             if (System.getProperty("os.name").startsWith("Win")) {
                 String path = Utils.getenv("PATH");
                 if (path == null) {
@@ -199,6 +274,33 @@ public final class CudaUtils {
         }
     }
 
+    private static String[] execute(int deviceId) {
+        try {
+            String javaHome = System.getProperty("java.home");
+            String classPath = System.getProperty("java.class.path");
+            String os = System.getProperty("os.name");
+            List<String> cmd = new ArrayList<>(4);
+            if (os.startsWith("Win")) {
+                cmd.add(javaHome + "\\bin\\java.exe");
+            } else {
+                cmd.add(javaHome + "/bin/java");
+            }
+            cmd.add("-cp");
+            cmd.add(classPath);
+            cmd.add("ai.djl.util.cuda.CudaUtils");
+            if (deviceId >= 0) {
+                cmd.add(String.valueOf(deviceId));
+            }
+            Process ps = new ProcessBuilder(cmd).redirectErrorStream(true).start();
+            try (InputStream is = ps.getInputStream()) {
+                String line = Utils.toString(is).trim();
+                return line.split(",");
+            }
+        } catch (IOException e) {
+            throw new IllegalArgumentException("Failed get GPU information", e);
+        }
+    }
+
     private static void checkCall(int ret) {
         if (LIB == null) {
             throw new IllegalStateException("No cuda library is loaded.");
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngineProvider.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngineProvider.java
index f42f691d2..ad440a479 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngineProvider.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngineProvider.java
@@ -18,6 +18,8 @@ import ai.djl.engine.EngineProvider;
 /** {@code TfEngineProvider} is the TensorFlow implementation of {@link EngineProvider}. */
 public class TfEngineProvider implements EngineProvider {
 
+    private static volatile Engine engine; // NOPMD
+
     /** {@inheritDoc} */
     @Override
     public String getEngineName() {
@@ -33,10 +35,13 @@ public class TfEngineProvider implements EngineProvider {
     /** {@inheritDoc} */
     @Override
     public Engine getEngine() {
-        return InstanceHolder.INSTANCE;
-    }
-
-    private static class InstanceHolder {
-        static final Engine INSTANCE = TfEngine.newInstance();
+        if (engine == null) {
+            synchronized (TfEngineProvider.class) {
+                if (engine == null) {
+                    engine = TfEngine.newInstance();
+                }
+            }
+        }
+        return engine;
     }
 }


Test patch:
diff --git a/api/src/test/java/ai/djl/util/SecurityManagerTest.java b/api/src/test/java/ai/djl/util/SecurityManagerTest.java
index fd9b5db72..1e9eb17f6 100644
--- a/api/src/test/java/ai/djl/util/SecurityManagerTest.java
+++ b/api/src/test/java/ai/djl/util/SecurityManagerTest.java
@@ -74,8 +74,11 @@ public class SecurityManagerTest {
                     }
                 };
         System.setSecurityManager(sm);
-
-        Assert.assertFalse(CudaUtils.hasCuda());
-        Assert.assertEquals(CudaUtils.getGpuCount(), 0);
+        try {
+            Assert.assertFalse(CudaUtils.hasCuda());
+            Assert.assertEquals(CudaUtils.getGpuCount(), 0);
+        } finally {
+            System.setSecurityManager(null);
+        }
     }
 }
diff --git a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
index de1c5cb4a..a598d8482 100644
--- a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
+++ b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
@@ -20,8 +20,6 @@ import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.lang.management.MemoryUsage;
-import java.util.Arrays;
-import java.util.List;
 
 public class CudaUtilsTest {
 
@@ -30,6 +28,9 @@ public class CudaUtilsTest {
     @Test
     public void testCudaUtils() {
         if (!CudaUtils.hasCuda()) {
+            Assert.assertThrows(CudaUtils::getCudaVersionString);
+            Assert.assertThrows(() -> CudaUtils.getComputeCapability(0));
+            Assert.assertThrows(() -> CudaUtils.getGpuMemory(Device.gpu()));
             return;
         }
         // Possible to have CUDA and not have a GPU.
@@ -37,16 +38,24 @@ public class CudaUtilsTest {
             return;
         }
 
-        int cudaVersion = CudaUtils.getCudaVersion();
+        String cudaVersion = CudaUtils.getCudaVersionString();
         String smVersion = CudaUtils.getComputeCapability(0);
         MemoryUsage memoryUsage = CudaUtils.getGpuMemory(Device.gpu());
 
         logger.info("CUDA runtime version: {}, sm: {}", cudaVersion, smVersion);
         logger.info("Memory usage: {}", memoryUsage);
 
-        Assert.assertTrue(cudaVersion >= 9020, "cuda 9.2+ required.");
+        Assert.assertNotNull(cudaVersion);
+        Assert.assertNotNull(smVersion);
+    }
 
-        List<String> supportedSm = Arrays.asList("37", "52", "60", "61", "70", "75");
-        Assert.assertTrue(supportedSm.contains(smVersion), "Unsupported cuda sm: " + smVersion);
+    @Test
+    public void testCudaUtilsWithFolk() {
+        System.setProperty("ai.djl.util.cuda.folk", "true");
+        try {
+            testCudaUtils();
+        } finally {
+            System.clearProperty("ai.djl.util.cuda.folk");
+        }
     }
 }