Cache the PyTorch "files.txt". (#1982)

* Cache the PyTorch "files.txt". The cached "files.txt" ensures DJL will still run when the machine is offline, and the native torch jar is not added to the classpath. Co-authored-by: Frank Liu <[email protected]>
deepjavalibrary · Sep 2, 2022 · ae05f94 · ae05f94
1 parent 0d8f745
commit ae05f94
Showing 1 changed file with 16 additions and 1 deletion.
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
@@ -387,7 +387,22 @@ private static LibTorch downloadPyTorch(Platform platform) {
         }
         String link = "https://publish.djl.ai/pytorch/" + matcher.group(1);
         Path tmp = null;
-        try (InputStream is = new URL(link + "/files.txt").openStream()) {
+
+        Path indexFile = cacheDir.resolve(version + ".txt");
+        if (Files.notExists(indexFile)) {
+            Path tempFile = cacheDir.resolve(version + ".tmp");
+            try (InputStream is = new URL(link + "/files.txt").openStream()) {
+                Files.createDirectories(cacheDir);
+                Files.copy(is, tempFile, StandardCopyOption.REPLACE_EXISTING);
+                Utils.moveQuietly(tempFile, indexFile);
+            } catch (IOException e) {
+                throw new IllegalStateException("Failed to save pytorch index file", e);
+            } finally {
+                Utils.deleteQuietly(tempFile);
+            }
+        }
+
+        try (InputStream is = Files.newInputStream(indexFile)) {
             // if files not found
             Files.createDirectories(cacheDir);
             List<String> lines = Utils.readLines(is);