[api] Workaround detect neuron issue on SageMaker (#2729)

deepjavalibrary · Jul 30, 2023 · d4ddc40 · d4ddc40
1 parent 03e2865
commit d4ddc40
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/api/src/main/java/ai/djl/util/NeuronUtils.java b/api/src/main/java/ai/djl/util/NeuronUtils.java
@@ -16,6 +16,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -86,6 +87,14 @@ public static int getNeuronCoresForDevice(String location) {
         if (!Files.exists(path)) {
             return 0;
         }
+        Path file = path.resolve("core_count");
+        if (Files.exists(file)) {
+            try (InputStream is = Files.newInputStream(file)) {
+                return Integer.parseInt(Utils.toString(is));
+            } catch (IOException e) {
+                throw new AssertionError("Failed to read core_count file", e);
+            }
+        }
         int count = 0;
         try (Stream<Path> dev = Files.list(path)) {
             return Math.toIntExact(dev.filter(p -> matches(p, "neuron_core")).count());