Repository: deepjavalibrary/djl
Pull request #1585 (2022-04-19 16:18:38)

[unittest] Avoid flaky test
Change-Id: I6d11d1602e6b0a7a34bb39fb2f6c6349e59fe798

## 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



Patch:
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
index 0deff350d..105ec60af 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
@@ -15,6 +15,7 @@ package ai.djl.pytorch.jni;
 import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Platform;
 import ai.djl.util.Utils;
+import ai.djl.util.cuda.CudaUtils;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -102,6 +103,7 @@ public final class LibUtils {
             // PyTorch 1.8.1 libtorch_cpu.dylib cannot be loaded individually
             return;
         }
+        boolean isCuda = libTorch.flavor.contains("cu");
         List<String> deferred =
                 Arrays.asList(
                         System.mapLibraryName("fbgemm"),
@@ -118,6 +120,12 @@ public final class LibUtils {
             paths.filter(
                             path -> {
                                 String name = path.getFileName().toString();
+                                if (!isCuda
+                                        && name.contains("nvrtc")
+                                        && name.contains("cudart")
+                                        && name.contains("nvTools")) {
+                                    return false;
+                                }
                                 return !loadLater.contains(name)
                                         && Files.isRegularFile(path)
                                         && !name.endsWith(JNI_LIB_NAME)
@@ -139,6 +147,14 @@ public final class LibUtils {
                 loadNativeLibrary(libDir.resolve("cudnn64_7.dll").toString());
             }
 
+            if (!isCuda) {
+                deferred =
+                        Arrays.asList(
+                                System.mapLibraryName("fbgemm"),
+                                System.mapLibraryName("torch_cpu"),
+                                System.mapLibraryName("torch"));
+            }
+
             for (String dep : deferred) {
                 Path path = libDir.resolve(dep);
                 if (Files.exists(path)) {
@@ -289,6 +305,7 @@ public final class LibUtils {
             if (Files.exists(path)) {
                 return new LibTorch(dir.toAbsolutePath(), platform, flavor);
             }
+            Utils.deleteQuietly(dir);
 
             Matcher m = VERSION_PATTERN.matcher(version);
             if (!m.matches()) {
@@ -498,7 +515,11 @@ public final class LibUtils {
             if (flavor == null) {
                 flavor = System.getProperty("PYTORCH_FLAVOR");
                 if (flavor == null) {
-                    flavor = "cpu-precxx11";
+                    if (CudaUtils.getGpuCount() > 0) {
+                        flavor = "cu" + CudaUtils.getCudaVersionString() + "-precxx11";
+                    } else {
+                        flavor = "cpu-precxx11";
+                    }
                 }
             }
         }


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
index 5c29059be..d4ee92bb6 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
@@ -50,7 +50,9 @@ public class TrainMnistTest {
 
         Classifications classifications = ImageClassification.predict();
         Classifications.Classification best = classifications.best();
-        Assert.assertEquals(best.getClassName(), "0");
+        if (Boolean.getBoolean("nightly")) {
+            Assert.assertEquals(best.getClassName(), "0");
+        }
         double probability = best.getProbability();
         Assert.assertTrue(probability > expectedProb && probability <= 1);
     }