use jars pre determined sizes

bioimage-io · Jan 8, 2025 · be6e46d · be6e46d
1 parent ea99953
commit be6e46d
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/bioimageio/download/DownloadTracker.java b/src/main/java/io/bioimage/modelrunner/bioimageio/download/DownloadTracker.java
@@ -34,6 +34,7 @@
 import java.util.stream.Collectors;
 
 import io.bioimage.modelrunner.engine.installation.EngineInstall;
+import io.bioimage.modelrunner.versionmanagement.JarInfo;
 
 /**
  * Class that contains the methods to track the progress downloading files. 
@@ -136,14 +137,15 @@ private DownloadTracker(String folder, TwoParameterConsumer<String, Double> cons
 		this.parentThread = Thread.currentThread();
 		this.folder = folder;
 		sizeFiles = new LinkedHashMap<String, Long>();
+		JarInfo jarInfo = JarInfo.getInstance();
 		for (String link : links) {
 			String key = folder + File.separator + DownloadModel.getFileNameFromURLString(link);
 			if (consumer.get().get(key + EngineInstall.NBYTES_SUFFIX) != null && consumer.get().get(key + EngineInstall.NBYTES_SUFFIX) != -1) {
 				sizeFiles.put(key, consumer.get().get(key + EngineInstall.NBYTES_SUFFIX).longValue());
 				continue;
 			}
 			try {
-				sizeFiles.put(key, DownloadModel.getFileSize(new URL(link)));
+				sizeFiles.put(key, (jarInfo.get(link) != null ? jarInfo.get(link) : DownloadModel.getFileSize(new URL(link))));
 			} catch (MalformedURLException e) {
 				throw new IOException("The URL '" + link + "' cannot be found.");
 			}

diff --git a/src/main/java/io/bioimage/modelrunner/engine/installation/EngineInstall.java b/src/main/java/io/bioimage/modelrunner/engine/installation/EngineInstall.java
@@ -54,6 +54,7 @@
 import io.bioimage.modelrunner.versionmanagement.AvailableEngines;
 import io.bioimage.modelrunner.versionmanagement.DeepLearningVersion;
 import io.bioimage.modelrunner.versionmanagement.InstalledEngines;
+import io.bioimage.modelrunner.versionmanagement.JarInfo;
 
 /**
  * Class that manages the dl-modelrunner engines.
@@ -596,13 +597,14 @@ public long getBasicDownloadTotalSize() throws InterruptedException {
 		long totalSize = 0;
 		for (Entry<String, String> ee : missingEngineFolders.entrySet()) {
 			try {
+				JarInfo jarInfo = JarInfo.getInstance();
 				long engineSize = 0;
 				DeepLearningVersion dlVersion = DeepLearningVersion.fromFile(new File(ee.getValue()));
 				for (String link : dlVersion.getJars()) {
 					if (Thread.currentThread().isInterrupted())
 						throw new InterruptedException("Retrieval of download size interrupted.");
 					String key = ee.getValue() + File.separator + DownloadModel.getFileNameFromURLString(link) + NBYTES_SUFFIX;
-					long val = DownloadModel.getFileSize(new URL(link));
+					long val = jarInfo.get(link) != null ? jarInfo.get(link) : DownloadModel.getFileSize(new URL(link));
 					this.consumersMap.get(ee.getValue()).accept(key, (double) val);
 					engineSize += val;
 				}

diff --git a/src/main/java/io/bioimage/modelrunner/versionmanagement/JarInfo.java b/src/main/java/io/bioimage/modelrunner/versionmanagement/JarInfo.java
@@ -39,7 +39,7 @@ public static JarInfo getInstance() throws IOException {
      *
      * @return Map containing the URL and their respective sizes
      */
-    public Map<String, Long> getUrlData() {
+    public Map<String, Long> getAllData() {
         return urlData;
     }
 
@@ -49,7 +49,7 @@ public Map<String, Long> getUrlData() {
      * @param url The URL to look up
      * @return The size associated with the URL, or null if not found
      */
-    public Long getSizeForUrl(String url) {
+    public Long get(String url) {
         return urlData.get(url);
     }