[pytorch] Allows to load libstdc++.so.6 form different location (#2929)

deepjavalibrary · Jan 9, 2024 · f1b9da6 · f1b9da6
1 parent 6141c48
commit f1b9da6
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
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
@@ -106,6 +106,17 @@ public static String getLibtorchPath() {
 
     private static void loadLibTorch(LibTorch libTorch) {
         Path libDir = libTorch.dir.toAbsolutePath();
+        if (Files.exists(libDir.resolve("libstdc++.so.6"))) {
+            String libstd = Utils.getEnvOrSystemProperty("LIBSTDCXX_LIBRARY_PATH");
+            if (libstd != null) {
+                try {
+                    logger.info("Loading libstdc++.so.6 from: {}", libstd);
+                    System.load(libstd);
+                } catch (UnsatisfiedLinkError e) {
+                    logger.warn("Failed Loading libstdc++.so.6 from: {}", libstd);
+                }
+            }
+        }
         boolean isCuda = libTorch.flavor.contains("cu");
         List<String> deferred =
                 Arrays.asList(

diff --git a/...djl/pytorch/integration/LibUtilsTest.java → ...jl/pytorch/integration/ALibUtilsTest.java b/...djl/pytorch/integration/LibUtilsTest.java → ...jl/pytorch/integration/ALibUtilsTest.java
@@ -18,17 +18,21 @@
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
-public class LibUtilsTest {
+// Ensure this test run first
+public class ALibUtilsTest {
 
     @BeforeClass
     public void setup() {
-        System.setProperty(
-                "ai.djl.pytorch.native_helper", "ai.djl.pytorch.integration.LibUtilsTest");
+        System.setProperty("ai.djl.pytorch.native_helper", ALibUtilsTest.class.getName());
+        System.setProperty("STDCXX_LIBRARY_PATH", "/usr/lib/non-exists");
+        System.setProperty("PYTORCH_PRECXX11", "true");
     }
 
     @AfterClass
     public void teardown() {
         System.clearProperty("ai.djl.pytorch.native_helper");
+        System.clearProperty("LIBSTDCXX_LIBRARY_PATH");
+        System.clearProperty("PYTORCH_PRECXX11");
     }
 
     @Test