[tvm] branch unity updated: Add the ability to differentiate between model loads from remote fetch v/s model loads from cache (#15357)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new 057e61575b Add the ability to differentiate between model loads from 
remote fetch v/s model loads from cache (#15357)
057e61575b is described below

commit 057e61575b537b0c926e2f5ac07a37160780b1e6
Author: narangkay <kanar...@gmail.com>
AuthorDate: Wed Jul 19 09:09:40 2023 -0700

    Add the ability to differentiate between model loads from remote fetch v/s 
model loads from cache (#15357)
    
    * add capability to report cache progress differently from full fetch 
progress
    
    * add helper methods to check cache
---
 web/src/index.ts   |  2 +-
 web/src/runtime.ts | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/web/src/index.ts b/web/src/index.ts
index 7d26fa7001..fd27fce9fd 100644
--- a/web/src/index.ts
+++ b/web/src/index.ts
@@ -22,7 +22,7 @@ export {
   PackedFunc, Module, NDArray,
   TVMArray, TVMObject, VirtualMachine,
   InitProgressCallback, InitProgressReport,
-  ArtifactCache, Instance, instantiate
+  ArtifactCache, Instance, instantiate, hasNDArrayInCache
 } from "./runtime";
 export { Disposable, LibraryProvider } from "./types";
 export { RPCServer } from "./rpc_server";
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 021bfceb1e..1ae3f1211b 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -972,6 +972,7 @@ export interface NDArrayShardEntry {
 export interface InitProgressReport {
   progress: number;
   timeElapsed: number;
+  cacheOnly: boolean;
   text: string;
 }
 
@@ -1003,6 +1004,16 @@ export class ArtifactCache {
     }
     return result;
   }
+
+  async hasAllKeys(keys: string[]) {
+    if (this.cache === undefined) {
+      this.cache = await caches.open(this.scope);
+    }
+    return this.cache.keys()
+      .then(requests => requests.map(request => request.url))
+      .then(cacheKeys => keys.every(key => cacheKeys.indexOf(key) !== -1))
+      .catch(err => false);
+  }
 }
 
 /**
@@ -1472,6 +1483,8 @@ export class Instance implements Disposable {
     let fetchedBytes = 0;
     let timeElapsed = 0;
 
+    const cacheOnly = await artifactCache.hasAllKeys(list.map(key => new 
URL(key.dataPath, ndarrayCacheUrl).href))
+
     const reportCallback = (iter: number) => {
       // report
       for (let j = 0; j < this.initProgressCallback.length; ++j) {
@@ -1481,9 +1494,16 @@ export class Instance implements Disposable {
         text += timeElapsed + " secs elapsed.";
         text += " It can take a while when we first visit this page to 
populate the cache."
         text += " Later refreshes will become faster.";
+        if (cacheOnly) {
+          text = "Loading model from cache[" + iter + "/" + list.length + "]: 
";
+          text += Math.ceil(fetchedBytes / (1024 * 1024)).toString() + "MB 
loaded. "
+          text += Math.floor(fetchedBytes * 100 / totalBytes).toString() + "% 
completed, "
+          text += timeElapsed + " secs elapsed.";
+        }
         this.initProgressCallback[j]({
           progress: fetchedBytes / totalBytes,
           timeElapsed: timeElapsed,
+          cacheOnly: cacheOnly,
           text: text
         });
       }
@@ -1493,6 +1513,7 @@ export class Instance implements Disposable {
       this.initProgressCallback[j]({
         progress: fetchedBytes / totalBytes,
         timeElapsed: 0,
+        cacheOnly: cacheOnly,
         text: "Start to fetch params",
       });
     }
@@ -1926,6 +1947,7 @@ export class Instance implements Disposable {
           this.initProgressCallback[j]({
             progress: progress,
             timeElapsed: timeElapsed,
+            cacheOnly: false,
             text: text
           });
         }
@@ -2349,3 +2371,22 @@ export function instantiate(
     }
   );
 }
+
+export async function hasNDArrayInCache(
+  ndarrayCacheUrl: string,
+  cacheScope: string = "tvmjs"
+): Promise<boolean> {
+  const artifactCache = new ArtifactCache(cacheScope);
+  const jsonUrl = new URL("ndarray-cache.json", ndarrayCacheUrl).href;
+  const hasJsonUrlInCache = await artifactCache.hasAllKeys([jsonUrl]);
+  if (!hasJsonUrlInCache) {
+    return false;
+  }
+  const result = await artifactCache.fetchWithCache(jsonUrl);
+  let list;
+  if (result instanceof Response) {
+    list = await result.json();
+  }
+  list = list["records"] as Array<NDArrayShardEntry>;
+  return await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, 
ndarrayCacheUrl).href));
+}