[api] Fixes nightly tests on GPU machine

api/src/test/java/ai/djl/util/SecurityManagerTest.java

@@ -12,14 +12,11 @@
  */
 package ai.djl.util;
 
-import ai.djl.util.cuda.CudaUtils;
-
 import org.testng.Assert;
 import org.testng.annotations.AfterTest;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
-import java.io.FilePermission;
 import java.security.Permission;
 
 public class SecurityManagerTest {
@@ -59,26 +56,4 @@ public void checkPermission(Permission perm) {
         Assert.assertEquals(Utils.getenv("TEST", "test"), "test");
         Assert.assertEquals(Utils.getenv().size(), 0);
     }
-
-    @Test
-    public void testCudaUtils() {
-        // Disable access to the cudart library files
-        SecurityManager sm =
-                new SecurityManager() {
-                    @Override
-                    public void checkPermission(Permission perm) {
-                        if (perm instanceof FilePermission && perm.getName().contains("cudart")) {
-                            throw new SecurityException(
-                                    "Don't have permission to read file: " + perm.getName());
-                        }
-                    }
-                };
-        System.setSecurityManager(sm);
-        try {
-            Assert.assertFalse(CudaUtils.hasCuda());
-            Assert.assertEquals(CudaUtils.getGpuCount(), 0);
-        } finally {
-            System.setSecurityManager(null);
-        }
-    }
 }

examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java

@@ -35,6 +35,7 @@ public class MaskDetectionTest {
     @Test
     public void testMaskDetection() throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
+        TestRequirements.notGpu();
 
         DetectedObjects result = MaskDetection.predict();
         logger.info("{}", result);

examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java

@@ -35,6 +35,7 @@ public void testObjectDetection() throws ModelException, TranslateException, IOE
         // throttling and will fail the test.
         TestRequirements.linux();
         TestRequirements.nightly();
+        TestRequirements.notGpu();
 
         DetectedObjects result = ObjectDetectionWithTensorflowSavedModel.predict();
 

examples/src/test/java/ai/djl/examples/inference/SuperResolutionTest.java

@@ -32,6 +32,7 @@ public class SuperResolutionTest {
     @Test
     public void testSuperResolution() throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
+        TestRequirements.notGpu();
 
         String imagePath = "src/test/resources/";
         Image fox = ImageFactory.getInstance().fromFile(Paths.get(imagePath + "fox.png"));

examples/src/test/java/ai/djl/examples/inference/UniversalSentenceEncoderTest.java

@@ -29,6 +29,7 @@ public class UniversalSentenceEncoderTest {
     public void testSentimentAnalysis() throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
         TestRequirements.nightly();
+        TestRequirements.notGpu();
 
         List<String> inputs = new ArrayList<>();
         inputs.add("The quick brown fox jumps over the lazy dog.");

examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java

@@ -28,7 +28,7 @@ public class Yolov8DetectionTest {
 
     @Test
     public void testYolov8Detection() throws ModelException, TranslateException, IOException {
-        TestRequirements.linux();
+        TestRequirements.notGpu();
 
         DetectedObjects result = Yolov8Detection.predict();
 

examples/src/test/java/ai/djl/examples/training/TrainResNetTest.java

@@ -32,6 +32,7 @@ public class TrainResNetTest {
     @Test
     public void testTrainResNet() throws ModelException, IOException, TranslateException {
         TestRequirements.nightly();
+        TestRequirements.gpu("PyTorch", 1);
 
         // Limit max 4 gpu for cifar10 training to make it converge faster.
         // and only train 10 batch for unit test.
@@ -47,7 +48,7 @@ public void testTrainResNetImperativeNightly()
             throws ModelException, IOException, TranslateException {
         TestRequirements.linux();
         TestRequirements.nightly();
-        TestRequirements.gpu("PyTorch");
+        TestRequirements.gpu("PyTorch", 4);
 
         // Limit max 4 gpu for cifar10 training to make it converge faster.
         // and only train 10 batch for unit test.

examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java

@@ -27,7 +27,7 @@ public void testTrainSentimentAnalysis()
             throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
         TestRequirements.nightly();
-        TestRequirements.gpu("MXNet");
+        TestRequirements.gpu("MXNet", 1);
 
         String[] args = {"-e", "1", "-g", "1", "--engine", "MXNet"};
         TrainSentimentAnalysis.runExample(args);

examples/src/test/java/ai/djl/testing/TestRequirements.java

@@ -15,6 +15,7 @@
 import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
 
+import ai.djl.util.cuda.CudaUtils;
 import org.testng.SkipException;
 
 import java.util.Arrays;
@@ -76,9 +77,16 @@ public static void linux() {
     }
 
     /** Requires a test have at least one gpu. */
-    public static void gpu(String engine) {
-        if (Engine.getEngine(engine).getGpuCount() == 0) {
-            throw new SkipException("This test requires a GPU to run");
+    public static void gpu(String engine, int numGpu) {
+        if (Engine.getEngine(engine).getGpuCount() < numGpu) {
+            throw new SkipException("This test requires " + numGpu + " GPUs to run");
+        }
+    }
+
+    /** Avoid OOM on GPUs with multiple engines. */
+    public static void notGpu() {
+        if (CudaUtils.getGpuCount() > 0) {
+            throw new SkipException("This test requires CPU only machine to run");
         }
     }
 }