Repository: deepjavalibrary/djl
Pull request #3301 (2024-07-04 23:51:05)

[rust] Fixes test failure on GPU
## 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: :warning: Please install the !['codecov app svg image'](https://github.com/codecov/engineering-team/assets/152432831/e90313f4-9d3a-4b63-8b54-cfe14e7ec20d) to ensure uploads and comments are reliably processed by Codecov.

## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3301?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 72.70%. Comparing base [(`725851b`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/725851b5ae512706899850cb3a5c50a76ba96984?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) to head [(`4dd1cb2`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/4dd1cb2853fdc0d6a8b0510a69a078f437905152?dropdown=coverage&el=desc&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    #3301      +/-   ##
============================================
- Coverage     72.71%   72.70%   -0.01%     
+ Complexity     7361     7359       -2     
============================================
  Files           677      677              
  Lines         32561    32561              
  Branches       3433     3433              
============================================
- Hits          23677    23675       -2     
- Misses         7258     7261       +3     
+ Partials       1626     1625       -1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3301?dropdown=coverage&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/Platform.java b/api/src/main/java/ai/djl/util/Platform.java
index 2ab25c2f3..dd31a3833 100644
--- a/api/src/main/java/ai/djl/util/Platform.java
+++ b/api/src/main/java/ai/djl/util/Platform.java
@@ -20,7 +20,9 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
+import java.util.ArrayList;
 import java.util.Enumeration;
+import java.util.List;
 import java.util.Properties;
 
 /**
@@ -71,32 +73,50 @@ public final class Platform {
         }
 
         Platform systemPlatform = Platform.fromSystem(engine);
-        Platform placeholder = null;
+        List<Platform> availablePlatforms = new ArrayList<>();
         while (urls.hasMoreElements()) {
             URL url = urls.nextElement();
             Platform platform = Platform.fromUrl(url);
             platform.apiVersion = systemPlatform.apiVersion;
             if (platform.isPlaceholder()) {
-                placeholder = platform;
+                availablePlatforms.add(platform);
             } else if (platform.matches(systemPlatform)) {
                 logger.info("Found matching platform from: {}", url);
-                return platform;
+                availablePlatforms.add(platform);
             } else {
                 logger.info("Ignore mismatching platform from: {}", url);
             }
         }
-        if (placeholder != null) {
-            logger.info("Found placeholder platform from: {}", placeholder);
-            return placeholder;
-        }
-
-        if (systemPlatform.version == null) {
-            throw new AssertionError("No " + engine + " version found in property file.");
-        }
-        if (systemPlatform.apiVersion == null) {
-            throw new AssertionError("No " + engine + " djl_version found in property file.");
+        if (availablePlatforms.isEmpty()) {
+            if (systemPlatform.version == null) {
+                throw new AssertionError("No " + engine + " version found in property file.");
+            }
+            if (systemPlatform.apiVersion == null) {
+                throw new AssertionError("No " + engine + " djl_version found in property file.");
+            }
+            return systemPlatform;
+        } else if (availablePlatforms.size() == 1) {
+            Platform ret = availablePlatforms.get(0);
+            if (ret.isPlaceholder()) {
+                logger.info("Found placeholder platform from: {}", ret);
+            }
+            return ret;
         }
-        return systemPlatform;
+        availablePlatforms.sort(
+                (o1, o2) -> {
+                    if (o1.isPlaceholder()) {
+                        return 1;
+                    } else if (o2.isPlaceholder()) {
+                        return -1;
+                    }
+                    // cu121-precx11 > cu121 > cu118-precss11 > cpu-precxx11 > cpu
+                    int ret = o2.getFlavor().compareTo(o1.getFlavor());
+                    if (ret == 0) {
+                        return o2.getVersion().compareTo(o1.getVersion());
+                    }
+                    return ret;
+                });
+        return availablePlatforms.get(0);
     }
 
     /**
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 859596a6b..b234eb479 100644
--- a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
+++ b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
@@ -139,11 +139,13 @@ public final class CudaUtils {
      */
     public static String getComputeCapability(int device) {
         if (Boolean.getBoolean("ai.djl.util.cuda.fork")) {
-            String[] ret = execute(device);
-            if (ret.length != 3) {
-                throw new IllegalArgumentException(ret[0]);
+            if (gpuInfo == null) { // NOPMD
+                gpuInfo = execute(-1);
             }
-            return ret[0];
+            if (device >= gpuInfo.length - 2) {
+                throw new IllegalArgumentException("Invalid device: " + device);
+            }
+            return gpuInfo[device + 2];
         }
 
         if (LIB == null) {
@@ -214,7 +216,12 @@ public final class CudaUtils {
                 return;
             }
             int cudaVersion = getCudaVersion();
-            System.out.println(gpuCount + "," + cudaVersion);
+            StringBuilder sb = new StringBuilder();
+            sb.append(gpuCount).append(',').append(cudaVersion);
+            for (int i = 0; i < gpuCount; ++i) {
+                sb.append(',').append(getComputeCapability(i));
+            }
+            System.out.println(sb);
             return;
         }
         try {
diff --git a/extensions/tokenizers/rust/src/layers/cublaslt.rs b/extensions/tokenizers/rust/src/layers/cublaslt.rs
index 0ecca7000..4fe5431e4 100644
--- a/extensions/tokenizers/rust/src/layers/cublaslt.rs
+++ b/extensions/tokenizers/rust/src/layers/cublaslt.rs
@@ -2,8 +2,8 @@
 pub use cudarc::cublaslt::Activation;
 
 use crate::layers::HiddenAct;
-use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
 use candle::backend::BackendStorage;
+use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
 
 #[cfg(feature = "cuda")]
 use candle::cuda_backend::WrapErr;
diff --git a/extensions/tokenizers/rust/src/layers/linear.rs b/extensions/tokenizers/rust/src/layers/linear.rs
index 966a672a6..2b67d777b 100644
--- a/extensions/tokenizers/rust/src/layers/linear.rs
+++ b/extensions/tokenizers/rust/src/layers/linear.rs
@@ -20,7 +20,12 @@ pub struct Linear {
 }
 
 impl Linear {
-    pub fn load(vb: VarBuilder, in_dim: usize, out_dim: usize, act: Option<HiddenAct>) -> Result<Self> {
+    pub fn load(
+        vb: VarBuilder,
+        in_dim: usize,
+        out_dim: usize,
+        act: Option<HiddenAct>,
+    ) -> Result<Self> {
         Ok(Self {
             weight: vb.get((out_dim, in_dim), "weight")?,
             bias: Some(vb.get(out_dim, "bias")?),
diff --git a/extensions/tokenizers/rust/src/layers/mod.rs b/extensions/tokenizers/rust/src/layers/mod.rs
index 378588fdf..3f0cd444c 100644
--- a/extensions/tokenizers/rust/src/layers/mod.rs
+++ b/extensions/tokenizers/rust/src/layers/mod.rs
@@ -2,5 +2,4 @@
 mod cublaslt;
 mod linear;
 
-pub use cublaslt::get_cublas_lt_wrapper;
 pub use linear::{HiddenAct, Linear};
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 5e37294a8..2d7a93e23 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -16,8 +16,8 @@ mod ndarray;
 
 #[cfg(feature = "cuda")]
 mod compute_cap;
-mod models;
 mod layers;
+mod models;
 
 extern crate tokenizers as tk;
 
@@ -30,7 +30,6 @@ use jni::errors::Error;
 use jni::objects::{
     JClass, JLongArray, JMethodID, JObject, JObjectArray, JString, JValue, ReleaseMode,
 };
-#[cfg(feature = "cuda")]
 use jni::sys::JNI_FALSE;
 
 use jni::sys::{jboolean, jint, jlong, jobjectArray, jsize, jvalue, JNI_TRUE};
@@ -811,7 +810,7 @@ pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_isCudaAvailable<'loca
     }
     #[cfg(not(feature = "cuda"))]
     {
-        JNI_TRUE
+        JNI_FALSE
     }
 }
 
diff --git a/extensions/tokenizers/rust/src/models/bert.rs b/extensions/tokenizers/rust/src/models/bert.rs
index 7d53e96fb..c241232d3 100644
--- a/extensions/tokenizers/rust/src/models/bert.rs
+++ b/extensions/tokenizers/rust/src/models/bert.rs
@@ -1,5 +1,5 @@
-use crate::models::Model;
 use crate::layers::Linear;
+use crate::models::Model;
 use candle::{Device, Result, Tensor};
 use candle_nn::{embedding, Embedding, Module, VarBuilder};
 use candle_transformers::models::with_tracing::{layer_norm, LayerNorm};
@@ -359,7 +359,12 @@ struct BertIntermediate {
 
 impl BertIntermediate {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
-        let dense = Linear::load(vb.pp("dense"), config.hidden_size, config.intermediate_size, None)?;
+        let dense = Linear::load(
+            vb.pp("dense"),
+            config.hidden_size,
+            config.intermediate_size,
+            None,
+        )?;
         Ok(Self {
             dense,
             intermediate_act: HiddenActLayer::new(config.hidden_act),
@@ -387,7 +392,12 @@ struct BertOutput {
 
 impl BertOutput {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
-        let dense = Linear::load(vb.pp("dense"), config.intermediate_size, config.hidden_size, None)?;
+        let dense = Linear::load(
+            vb.pp("dense"),
+            config.intermediate_size,
+            config.hidden_size,
+            None,
+        )?;
         let layer_norm = layer_norm(
             config.hidden_size,
             config.layer_norm_eps,
diff --git a/extensions/tokenizers/rust/src/models/mod.rs b/extensions/tokenizers/rust/src/models/mod.rs
index b82f00670..a0b52daa4 100644
--- a/extensions/tokenizers/rust/src/models/mod.rs
+++ b/extensions/tokenizers/rust/src/models/mod.rs
@@ -4,7 +4,7 @@ mod distilbert;
 use crate::ndarray::as_data_type;
 use crate::{cast_handle, drop_handle, to_handle, to_string_array};
 use bert::{BertConfig, BertModel};
-use candle::{Device, DType, Result, Tensor};
+use candle::{DType, Device, Result, Tensor};
 use candle_nn::VarBuilder;
 use distilbert::{DistilBertConfig, DistilBertModel};
 use jni::objects::{JLongArray, JObject, JString, ReleaseMode};


Test patch:
diff --git a/api/src/test/java/ai/djl/util/PlatformTest.java b/api/src/test/java/ai/djl/util/PlatformTest.java
index 039cbbea8..3e20609a0 100644
--- a/api/src/test/java/ai/djl/util/PlatformTest.java
+++ b/api/src/test/java/ai/djl/util/PlatformTest.java
@@ -12,16 +12,23 @@
  */
 package ai.djl.util;
 
+import ai.djl.util.cuda.CudaUtils;
+
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.Test;
 
 import java.io.BufferedWriter;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
 
 public class PlatformTest {
 
@@ -94,6 +101,43 @@ public class PlatformTest {
         Assert.assertFalse(platform.matches(system));
     }
 
+    @Test
+    public void testDetectPlatform() throws IOException, ReflectiveOperationException {
+        Path dir = Paths.get("build/tmp/");
+        Files.createDirectories(dir);
+        Platform system = Platform.fromSystem();
+        String classifier = system.getClassifier();
+        createZipFile(0, "1.0", "cpu", classifier, true);
+        createZipFile(1, "1.0", "cpu", classifier, false);
+        createZipFile(2, "1.0", "cpu-precxx11", classifier, false);
+        createZipFile(3, "1.0", "cu117", classifier, false);
+        createZipFile(4, "1.0", "cu117-precxx11", classifier, false);
+        createZipFile(5, "1.0", "cu999", classifier, false);
+        createZipFile(6, "1.0", "cu999-precxx11", classifier, false);
+        createZipFile(7, "99.99", "cu999-precxx11", classifier, false);
+        System.setProperty("ai.djl.util.cuda.fork", "true");
+        try {
+            String[] gpuInfo = new String[] {"1", "99990", "90"};
+            Field field = CudaUtils.class.getDeclaredField("gpuInfo");
+            field.setAccessible(true);
+            field.set(null, gpuInfo);
+            URL[] urls = new URL[8];
+            for (int i = 0; i < 8; ++i) {
+                urls[i] = dir.resolve(i + ".jar").toUri().toURL();
+            }
+            URLClassLoader cl = new URLClassLoader(urls);
+            Thread.currentThread().setContextClassLoader(cl);
+
+            Platform detected = Platform.detectPlatform("pytorch");
+            Assert.assertEquals(detected.getFlavor(), "cu999-precxx11");
+
+            field.set(null, null);
+        } finally {
+            System.clearProperty("ai.djl.util.cuda.fork");
+            Thread.currentThread().setContextClassLoader(null);
+        }
+    }
+
     private URL createPropertyFile(String content) throws IOException {
         Path dir = Paths.get("build/tmp/testFile/");
         Files.createDirectories(dir);
@@ -104,4 +148,24 @@ public class PlatformTest {
         }
         return file.toUri().toURL();
     }
+
+    private void createZipFile(
+            int index, String version, String flavor, String classifier, boolean placeHolder)
+            throws IOException {
+        Path file = Paths.get("build/tmp/" + index + ".jar");
+        try (JarOutputStream jos = new JarOutputStream(Files.newOutputStream(file))) {
+            JarEntry entry = new JarEntry("native/lib/pytorch.properties");
+            jos.putNextEntry(entry);
+            if (placeHolder) {
+                jos.write("placeholder=true\nversion=2.3.1".getBytes(StandardCharsets.UTF_8));
+            } else {
+                jos.write("version=".getBytes(StandardCharsets.UTF_8));
+                jos.write(version.getBytes(StandardCharsets.UTF_8));
+                jos.write("\nflavor=".getBytes(StandardCharsets.UTF_8));
+                jos.write(flavor.getBytes(StandardCharsets.UTF_8));
+                jos.write("\nclassifier=".getBytes(StandardCharsets.UTF_8));
+                jos.write(classifier.getBytes(StandardCharsets.UTF_8));
+            }
+        }
+    }
 }
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 df135938a..f6ecc66ac 100644
--- a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
+++ b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
@@ -55,6 +55,8 @@ public class CudaUtilsTest {
         System.setProperty("ai.djl.util.cuda.fork", "true");
         try {
             testCudaUtils();
+            CudaUtils.main(new String[0]);
+            CudaUtils.main(new String[] {"-1"});
         } finally {
             System.clearProperty("ai.djl.util.cuda.fork");
         }
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
index 4b8f883ff..3dddf301f 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.huggingface.tokenizers;
 
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
@@ -104,6 +105,7 @@ public class TextEmbeddingTranslatorTest {
                         .optArgument("block_shapes", "(1,7,384)")
                         .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
+                        .optDevice(Device.cpu())
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "max")
                         .optOption("hasParameter", "false")