From 1ed32239c9dd66ba07a9eed46ee8e7df4bf2fab3 Mon Sep 17 00:00:00 2001 From: DiegoCao Date: Mon, 18 Mar 2024 02:58:59 -0400 Subject: [PATCH 1/6] Support IndexDB for Larger model, modify artifact cache template --- web/src/artifact_cache.ts | 20 ++- web/src/index.ts | 2 +- web/src/runtime.ts | 357 ++++++++++++++++++++++++++++++++++---- 3 files changed, 335 insertions(+), 44 deletions(-) diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts index da9aaddfb0d6..616e36119317 100644 --- a/web/src/artifact_cache.ts +++ b/web/src/artifact_cache.ts @@ -21,15 +21,25 @@ */ export interface ArtifactCacheTemplate { /** - * fetch key url from cache + * fetch key url from cache, optional storetype for IndexDB + * + * storagetype for indexDB have two options: + * 1. json: return a json object + * 2. arraybuffer: return an arraybuffer object */ - fetchWithCache(url: string); + fetchWithCache(url: string, storetype?: string); /** - * add ey url to cache + * add key url to cache, optional storetype for IndexDB + * + * storagetype for indexDB have two options: + * 1. json: return a json object + * 2. arraybuffer: return an arraybuffer object + * + * returns the response or the specified stored object + * for reduced database transaction */ - addToCache(url: string); - + addToCache(url: string, storetype?: string): Promise; /** * check if cache has all keys in Cache */ diff --git a/web/src/index.ts b/web/src/index.ts index edc695978f50..d59dba3b46d7 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, hasNDArrayInCache, deleteNDArrayCache + ArtifactCache, ArtifactindexDBCache, Instance, instantiate, hasNDArrayInCache, deleteNDArrayCache } 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 9142571b9e4a..c93766888e99 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -1006,33 +1006,71 @@ export class ArtifactCache implements ArtifactCacheTemplate { this.scope = scope; } - async fetchWithCache(url: string) { + /** + * Convert the Response object to the expected storetype instead + * @param response the cache or indexDB response + * @param storetype the storetype stored in the indexedDB database + * @returns the expected response object + */ + async responseTostoretype(response: Response, storetype: string): Promise{ + let result: any; + if (storetype.toLowerCase() === "json"){ + result = await result.json(); + } else if (storetype.toLowerCase() === "arraybuffer") { + result = await result.arrayBuffer(); + } else { + console.error("Unknown storage type, return raw response"); + return response; + } + return result; + } + + /** + * fetch the corresponding url object in response or stored object format + * @param url url + * @param storetype the storage type for indexDB + * @returns response in json, arraybuffer or pure response format + */ + async fetchWithCache(url: string, storetype?: string) { const request = new Request(url); if (this.cache === undefined) { this.cache = await caches.open(this.scope); } let result = await this.cache.match(request); if (result === undefined) { - await this.cache.add(request); - result = await this.cache.match(request); - } - if (result === undefined) { - throw Error("Cannot fetch " + url); + result = await this.addToCache(url, storetype); + return result; + } else { + if (storetype === undefined){ + return result; + } else { + return await this.responseTostoretype(result, storetype); + } } - return result; } - async addToCache(url: string) { + async addToCache(url: string, storetype?: string) { const request = new Request(url); if (this.cache === undefined) { this.cache = await caches.open(this.scope); } - const result = await this.cache.match(request); + let result = await this.cache.match(request); if (result === undefined) { await this.cache.add(request); + result = await this.cache.match(request); + } + if (storetype === undefined){ + return result; + } else { + return await this.responseTostoretype(result, storetype); } } + /** + * Determine if all keys exist in the cache + * @param keys the url key list of the strings + * @returns boolean value indicate if all keys are in cache + */ async hasAllKeys(keys: string[]) { if (this.cache === undefined) { this.cache = await caches.open(this.scope); @@ -1040,15 +1078,210 @@ export class ArtifactCache implements ArtifactCacheTemplate { return this.cache.keys() .then(requests => requests.map(request => request.url)) .then(cacheKeys => keys.every(key => cacheKeys.indexOf(key) !== -1)) - .catch(err => false); + .catch(() => false); } + /** + * Delete the corresponding url object in cache + * @param url the corresponding url object to be deleted + */ async deleteInCache(url: string) { if (this.cache === undefined) { this.cache = await caches.open(this.scope); } - const result = await this.cache.delete(url); - return result; + await this.cache.delete(url); + } +} + +/** + * Cache by IndexDB to support caching model data + */ +export class ArtifactindexDBCache implements ArtifactCacheTemplate { + private dbName?: string; + private dbVersion = 1; + private db: IDBDatabase | undefined; + + constructor(dbName: string){ + this.dbName = dbName; + } + + /** + * Init the indexed DB database if it is not initialized. + */ + private async initDB() { + if (this.db != null){ + return; // the db is already inialized + } + return new Promise((resolve, reject) => { + const request = indexedDB.open(this.dbName, this.dbVersion); + request.onupgradeneeded = (event) => { + this.db = (event.target as IDBOpenDBRequest).result; + if (!this.db.objectStoreNames.contains('urls')) { + this.db.createObjectStore('urls', { keyPath: 'url' }); + } + }; + request.onsuccess = (event) => { + this.db = (event.target as IDBOpenDBRequest).result; + resolve(); + }; + request.onerror = (event) => { + console.error("Database error: ", (event.target as IDBOpenDBRequest).error); + reject((event.target as IDBOpenDBRequest).error); + }; + }); + } + + /** + * Check if current url object is in indexedDB or not + * @param url the url link + * @returns boolean indicate if url object in indexedDB + */ + private async isUrlInDB(url: string): Promise { + return new Promise((resolve, reject) => { + const transaction = this.db?.transaction(['urls'], 'readonly'); + if (transaction === undefined){ + return false; + } + const store = transaction.objectStore('urls'); + const request = store.get(url); + request.onsuccess = () => { + resolve(request.result !== undefined); + }; + request.onerror = (event) => { + reject((event.target as IDBRequest).error); + }; + }); + } + + async asyncGetHelper(url: string){ + return new Promise((resolve, reject) => { + let result: any; + const transaction = this.db?.transaction(['urls'], 'readonly'); + if (transaction === undefined){ + return false; + } + transaction.oncomplete = () => resolve(result); + transaction.onerror = () => reject(transaction.error); + const objectStore = transaction.objectStore('urls'); + const getRequest = objectStore.get(url); + getRequest.onsuccess = () => { + result = getRequest.result; + } + }) + } + + async fetchWithCache(url: string, storetype?: string) { + await this.initDB(); // await the initDB process + const isInDB = await this.isUrlInDB(url); + if (!isInDB) { + const response = await this.addToCache(url, storetype); + return response; + } else { + // URL found in DB, just fetch without storing + const result = await this.asyncGetHelper(url); + if (result != null && typeof result === "object" && "data" in result){ + return result.data; + } else if (result === null){ + // previously null data in cache! + await this.deleteInCache(url); + const response = await this.addToCache(url, storetype); + return response; + } + return null; + } + } + + async addToIndexDB(url: string, response: any, storetype?: string){ + await this.initDB(); + let data: any; + if (storetype != undefined){ + if (storetype.toLowerCase() === "json"){ + data = await response.json(); + } else if (storetype.toLocaleLowerCase() === "arraybuffer"){ + data = await response.arrayBuffer(); + } else { + console.error("Unsupported Type in IndexDB"); + } + } + return new Promise((resolve, reject) => { + const transaction = this.db?.transaction(['urls'], 'readwrite'); + if (transaction === undefined){ + return; + } + const store = transaction.objectStore('urls'); + const request = store.add({data, url}); // Index DB follows a {value, key} format, instead of {key, value} format! + request.onsuccess = () => resolve(); + request.onerror = (event) => reject((event.target as IDBRequest).error); + }); + } + + async addToCache(url: string, storetype?: string) :Promise{ + let response: Response; + try { + response = await fetch(url); + if (!response.ok) { + throw new Error('Network response was not ok'); + } + const response_copy = response.clone(); + await this.addToIndexDB(url, response_copy, storetype); + if (storetype.toLowerCase() === "arraybuffer"){ + return await response.arrayBuffer(); + } else if (storetype.toLowerCase() === "json"){ + return await response.json(); + } else { + return response; + } + } catch (error) { + console.error("There was a problem fetching the data:", error); + } + } + + async hasAllKeys(keys: string[]) :Promise { + await this.initDB(); // Ensure the DB is initialized + if (!this.db) { + throw new Error('Database is not initialized'); + } + return new Promise((resolve, reject) => { + const transaction = this.db.transaction(['urls'], 'readonly'); + const store = transaction.objectStore('urls'); + const promises = keys.map(key => { + return new Promise((resolve) => { + const request = store.get(key); + request.onsuccess = () => { + if (request.result === undefined) { + resolve(false); // Key not found, resolve with false + } else { + resolve(true); // Key found, resolve with true + } + }; + request.onerror = () => { + resolve(false); // On error, resolve as if the key was not found + }; + }); + }); + Promise.all(promises).then(results => { + const allExist = results.every(exists => exists); + resolve(allExist); + }).catch(error => { + reject(error); // Reject the main promise if any of the promises are rejected + }); + }); + } + + async deleteInCache(url: string) { + await this.initDB(); // Make sure the DB is initialized + const transaction = this.db?.transaction(['urls'], 'readwrite'); + if (transaction === undefined){ + return; + } + const store = transaction.objectStore('urls'); + const request = store.delete(url); + // Await completion of the delete request + await new Promise((resolve, reject) => { + request.onsuccess = () => resolve(); + request.onerror = () => reject(request.error); + }); + return; } } @@ -1500,20 +1733,34 @@ export class Instance implements Disposable { * @param ndarrayCacheUrl The cache url. * @param device The device to be fetched to. * @param cacheScope The scope identifier of the cache + * @param cacheType The type of the cache: "cache" or "indexDB" * @returns The meta data */ async fetchNDArrayCache( ndarrayCacheUrl: string, device: DLDevice, - cacheScope = "tvmjs" + cacheScope = "tvmjs", + cacheType = "cache" ): Promise { - const artifactCache = new ArtifactCache(cacheScope); + let artifactCache; + if (cacheType === undefined){ + artifactCache = new ArtifactCache(cacheScope); + } + if (cacheType.toLowerCase() === "cache"){ + artifactCache = new ArtifactCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexdb"){ + artifactCache = new ArtifactindexDBCache(cacheScope); + } else { + console.error("Unsupported Cache Type, using default browser cache"); + artifactCache = new ArtifactCache(cacheScope); + } const jsonUrl = new URL("ndarray-cache.json", ndarrayCacheUrl).href; - const result = await artifactCache.fetchWithCache(jsonUrl); - + const result = await artifactCache.fetchWithCache(jsonUrl, "json"); let list; if (result instanceof Response) { list = await result.json(); + } else { + list = result; } await this.fetchNDArrayCacheInternal( ndarrayCacheUrl, @@ -1538,7 +1785,6 @@ export class Instance implements Disposable { ) { const perf = compact.getPerformance(); const tstart = perf.now(); - let totalBytes = 0; for (let i = 0; i < list.length; ++i) { totalBytes += list[i].nbytes; @@ -1547,8 +1793,7 @@ export class Instance implements Disposable { let fetchedShards = 0; let timeElapsed = 0; - const cacheOnly = await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)) - + const cacheOnly = await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)); const reportCallback = (iter: number, loading = false) => { // report for (let j = 0; j < this.initProgressCallback.length; ++j) { @@ -1593,7 +1838,7 @@ export class Instance implements Disposable { const shard = list[i]; const dataUrl = new URL(shard.dataPath, ndarrayCacheUrl).href; try { - await artifactCache.addToCache(dataUrl); + await artifactCache.addToCache(dataUrl, "arraybuffer"); } catch (err) { this.env.logger("Error: Cannot fetch " + dataUrl + " err= " + err); throw err; @@ -1604,14 +1849,16 @@ export class Instance implements Disposable { } } // We launch 4 parallel for loops to limit the max concurrency to 4 download - const loopSize = Math.floor(list.length / 4); - await Promise.all([ - downloadCache(0, loopSize), - downloadCache(loopSize, 2 * loopSize), - downloadCache(2 * loopSize, 3 * loopSize), - downloadCache(3 * loopSize, list.length) - ]); - reportCallback(list.length, /*loading=*/true); + if (!cacheOnly){ + const loopSize = Math.floor(list.length / 4); + await Promise.all([ + downloadCache(0, loopSize), + downloadCache(loopSize, 2 * loopSize), + downloadCache(2 * loopSize, 3 * loopSize), + downloadCache(3 * loopSize, list.length) + ]); + reportCallback(list.length, /*loading=*/true); + } // Then iteratively, load the shard from cache for (let i = 0; i < list.length; ++i) { @@ -1619,7 +1866,7 @@ export class Instance implements Disposable { const dataUrl = new URL(shard.dataPath, ndarrayCacheUrl).href; let buffer; try { - buffer = await (await artifactCache.fetchWithCache(dataUrl)).arrayBuffer(); + buffer = await artifactCache.fetchWithCache(dataUrl, "arraybuffer"); } catch (err) { this.env.logger("Error: Cannot fetch " + dataUrl + " err= " + err); throw err; @@ -1661,6 +1908,9 @@ export class Instance implements Disposable { throw err; } } + if (cacheOnly){ + reportCallback(i + 1, /* Need to Report call back*/false); + } } } @@ -2118,7 +2368,6 @@ export class Instance implements Disposable { }).then(() => { finishCounter += 1; const tend = perf.now(); - const timeReportGap = 1000; // skip report if gap is smaller than 1000 if ((tend - tlastReport) < 1000 && finishCounter != fmapEntries.length) { return; @@ -2584,41 +2833,73 @@ export function instantiate( ); } +/** + * Function to check if NDarray is in Cache or not + * + * @param ndarrayCacheUrl The cache url which links to the NDArray + * @param cacheScope The scope identifier of the cache + * @param cacheType The type of the cache: "cache" or "indexDB" + * @returns the result if the cache has NDArray + */ export async function hasNDArrayInCache( ndarrayCacheUrl: string, - cacheScope = "tvmjs" + cacheScope = "tvmjs", + cacheType = "cache" ): Promise { - const artifactCache = new ArtifactCache(cacheScope); + let artifactCache; + if (cacheType.toLowerCase() === "cache"){ + artifactCache = new ArtifactCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexdb"){ + artifactCache = new ArtifactindexDBCache(cacheScope); + } else { + console.error("Unsupported Cache Type, using default browser cache"); + 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); + const result = await artifactCache.fetchWithCache(jsonUrl, "json"); let list; if (result instanceof Response) { list = await result.json(); + } else { + list = result; } list = list["records"] as Array; return await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)); } + /** * Given cacheUrl, search up items to delete based on cacheUrl/ndarray-cache.json * - * @param cacheUrl - * @param cacheScope + * @param cacheUrl The cacheUrl for the items + * @param cacheScope The scope identifier of the cache + * @param cacheType The type of the cache: "cache" or "indexDB" */ export async function deleteNDArrayCache( cacheUrl: string, - cacheScope = "tvmjs" + cacheScope = "tvmjs", + cacheType = "cache" ) { - const artifactCache = new ArtifactCache(cacheScope); + let artifactCache; + if (cacheType.toLowerCase() === "cache"){ + artifactCache = new ArtifactCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexdb"){ + artifactCache = new ArtifactindexDBCache(cacheScope); + } else { + console.error("Unsupported Cache Type, using default browser cache"); + artifactCache = new ArtifactCache(cacheScope); + } const jsonUrl = new URL("ndarray-cache.json", cacheUrl).href; - const result = await artifactCache.fetchWithCache(jsonUrl); + const result = await artifactCache.fetchWithCache(jsonUrl, "json"); let list; if (result instanceof Response) { list = await result.json(); + } else { + list = result; } const arrayentry = list["records"] as Array; const processShard = async (i: number) => { From 2f1f06e0f0d07e300055085d535c1628d45257fb Mon Sep 17 00:00:00 2001 From: Charlie Ruan <53290280+CharlieFRuan@users.noreply.github.com> Date: Fri, 5 Apr 2024 14:44:32 -0400 Subject: [PATCH 2/6] Minor formatting --- web/src/runtime.ts | 62 +++++++++++++++++++++++----------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/web/src/runtime.ts b/web/src/runtime.ts index c93766888e99..983e2fffa1d6 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -1012,9 +1012,9 @@ export class ArtifactCache implements ArtifactCacheTemplate { * @param storetype the storetype stored in the indexedDB database * @returns the expected response object */ - async responseTostoretype(response: Response, storetype: string): Promise{ + async responseTostoretype(response: Response, storetype: string): Promise { let result: any; - if (storetype.toLowerCase() === "json"){ + if (storetype.toLowerCase() === "json") { result = await result.json(); } else if (storetype.toLowerCase() === "arraybuffer") { result = await result.arrayBuffer(); @@ -1041,7 +1041,7 @@ export class ArtifactCache implements ArtifactCacheTemplate { result = await this.addToCache(url, storetype); return result; } else { - if (storetype === undefined){ + if (storetype === undefined) { return result; } else { return await this.responseTostoretype(result, storetype); @@ -1059,7 +1059,7 @@ export class ArtifactCache implements ArtifactCacheTemplate { await this.cache.add(request); result = await this.cache.match(request); } - if (storetype === undefined){ + if (storetype === undefined) { return result; } else { return await this.responseTostoretype(result, storetype); @@ -1101,7 +1101,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { private dbVersion = 1; private db: IDBDatabase | undefined; - constructor(dbName: string){ + constructor(dbName: string) { this.dbName = dbName; } @@ -1109,7 +1109,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { * Init the indexed DB database if it is not initialized. */ private async initDB() { - if (this.db != null){ + if (this.db != null) { return; // the db is already inialized } return new Promise((resolve, reject) => { @@ -1139,7 +1139,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { private async isUrlInDB(url: string): Promise { return new Promise((resolve, reject) => { const transaction = this.db?.transaction(['urls'], 'readonly'); - if (transaction === undefined){ + if (transaction === undefined) { return false; } const store = transaction.objectStore('urls'); @@ -1153,11 +1153,11 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { }); } - async asyncGetHelper(url: string){ + async asyncGetHelper(url: string) { return new Promise((resolve, reject) => { let result: any; const transaction = this.db?.transaction(['urls'], 'readonly'); - if (transaction === undefined){ + if (transaction === undefined) { return false; } transaction.oncomplete = () => resolve(result); @@ -1179,9 +1179,9 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } else { // URL found in DB, just fetch without storing const result = await this.asyncGetHelper(url); - if (result != null && typeof result === "object" && "data" in result){ + if (result != null && typeof result === "object" && "data" in result) { return result.data; - } else if (result === null){ + } else if (result === null) { // previously null data in cache! await this.deleteInCache(url); const response = await this.addToCache(url, storetype); @@ -1191,13 +1191,13 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } } - async addToIndexDB(url: string, response: any, storetype?: string){ + async addToIndexDB(url: string, response: any, storetype?: string) { await this.initDB(); let data: any; - if (storetype != undefined){ - if (storetype.toLowerCase() === "json"){ + if (storetype != undefined) { + if (storetype.toLowerCase() === "json") { data = await response.json(); - } else if (storetype.toLocaleLowerCase() === "arraybuffer"){ + } else if (storetype.toLocaleLowerCase() === "arraybuffer") { data = await response.arrayBuffer(); } else { console.error("Unsupported Type in IndexDB"); @@ -1205,17 +1205,17 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } return new Promise((resolve, reject) => { const transaction = this.db?.transaction(['urls'], 'readwrite'); - if (transaction === undefined){ + if (transaction === undefined) { return; } const store = transaction.objectStore('urls'); - const request = store.add({data, url}); // Index DB follows a {value, key} format, instead of {key, value} format! + const request = store.add({ data, url }); // Index DB follows a {value, key} format, instead of {key, value} format! request.onsuccess = () => resolve(); request.onerror = (event) => reject((event.target as IDBRequest).error); }); } - async addToCache(url: string, storetype?: string) :Promise{ + async addToCache(url: string, storetype?: string): Promise { let response: Response; try { response = await fetch(url); @@ -1224,9 +1224,9 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } const response_copy = response.clone(); await this.addToIndexDB(url, response_copy, storetype); - if (storetype.toLowerCase() === "arraybuffer"){ + if (storetype.toLowerCase() === "arraybuffer") { return await response.arrayBuffer(); - } else if (storetype.toLowerCase() === "json"){ + } else if (storetype.toLowerCase() === "json") { return await response.json(); } else { return response; @@ -1236,7 +1236,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } } - async hasAllKeys(keys: string[]) :Promise { + async hasAllKeys(keys: string[]): Promise { await this.initDB(); // Ensure the DB is initialized if (!this.db) { throw new Error('Database is not initialized'); @@ -1271,7 +1271,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { async deleteInCache(url: string) { await this.initDB(); // Make sure the DB is initialized const transaction = this.db?.transaction(['urls'], 'readwrite'); - if (transaction === undefined){ + if (transaction === undefined) { return; } const store = transaction.objectStore('urls'); @@ -1743,12 +1743,12 @@ export class Instance implements Disposable { cacheType = "cache" ): Promise { let artifactCache; - if (cacheType === undefined){ + if (cacheType === undefined) { artifactCache = new ArtifactCache(cacheScope); } - if (cacheType.toLowerCase() === "cache"){ + if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexdb"){ + } else if (cacheType.toLowerCase() == "indexdb") { artifactCache = new ArtifactindexDBCache(cacheScope); } else { console.error("Unsupported Cache Type, using default browser cache"); @@ -1849,7 +1849,7 @@ export class Instance implements Disposable { } } // We launch 4 parallel for loops to limit the max concurrency to 4 download - if (!cacheOnly){ + if (!cacheOnly) { const loopSize = Math.floor(list.length / 4); await Promise.all([ downloadCache(0, loopSize), @@ -1908,7 +1908,7 @@ export class Instance implements Disposable { throw err; } } - if (cacheOnly){ + if (cacheOnly) { reportCallback(i + 1, /* Need to Report call back*/false); } } @@ -2847,9 +2847,9 @@ export async function hasNDArrayInCache( cacheType = "cache" ): Promise { let artifactCache; - if (cacheType.toLowerCase() === "cache"){ + if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexdb"){ + } else if (cacheType.toLowerCase() == "indexdb") { artifactCache = new ArtifactindexDBCache(cacheScope); } else { console.error("Unsupported Cache Type, using default browser cache"); @@ -2885,9 +2885,9 @@ export async function deleteNDArrayCache( cacheType = "cache" ) { let artifactCache; - if (cacheType.toLowerCase() === "cache"){ + if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexdb"){ + } else if (cacheType.toLowerCase() == "indexdb") { artifactCache = new ArtifactindexDBCache(cacheScope); } else { console.error("Unsupported Cache Type, using default browser cache"); From 0a601925c1ba11d8e9c48bda26191fe996d14f12 Mon Sep 17 00:00:00 2001 From: Charlie Ruan <53290280+CharlieFRuan@users.noreply.github.com> Date: Fri, 5 Apr 2024 17:47:50 -0400 Subject: [PATCH 3/6] Rename indexdb to indexeddb --- web/src/artifact_cache.ts | 10 +++++----- web/src/index.ts | 5 ++--- web/src/runtime.ts | 34 +++++++++++++++++----------------- 3 files changed, 24 insertions(+), 25 deletions(-) diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts index 616e36119317..f1bf55921941 100644 --- a/web/src/artifact_cache.ts +++ b/web/src/artifact_cache.ts @@ -21,18 +21,18 @@ */ export interface ArtifactCacheTemplate { /** - * fetch key url from cache, optional storetype for IndexDB + * fetch key url from cache, optional storetype for IndexedDB * - * storagetype for indexDB have two options: - * 1. json: return a json object + * storagetype for indexedDB have two options: + * @param url: return a json object * 2. arraybuffer: return an arraybuffer object */ fetchWithCache(url: string, storetype?: string); /** - * add key url to cache, optional storetype for IndexDB + * add key url to cache, optional storetype for IndexedDB * - * storagetype for indexDB have two options: + * storagetype for indexedDB have two options: * 1. json: return a json object * 2. arraybuffer: return an arraybuffer object * diff --git a/web/src/index.ts b/web/src/index.ts index d59dba3b46d7..f134b0b5dc0b 100644 --- a/web/src/index.ts +++ b/web/src/index.ts @@ -22,11 +22,10 @@ export { PackedFunc, Module, NDArray, TVMArray, TVMObject, VirtualMachine, InitProgressCallback, InitProgressReport, - ArtifactCache, ArtifactindexDBCache, Instance, instantiate, hasNDArrayInCache, deleteNDArrayCache + ArtifactCache, ArtifactIndexedDBCache, Instance, instantiate, hasNDArrayInCache, deleteNDArrayCache } from "./runtime"; export { Disposable, LibraryProvider } from "./types"; export { RPCServer } from "./rpc_server"; -export { wasmPath, LinearCongruentialGenerator } from "./support"; +export { assert, wasmPath, LinearCongruentialGenerator } from "./support"; export { detectGPUDevice, GPUDeviceDetectOutput } from "./webgpu"; -export { assert } from "./support"; export { createPolyfillWASI } from "./compact"; diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 983e2fffa1d6..0d15520dfce3 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -996,7 +996,7 @@ export interface InitProgressReport { export type InitProgressCallback = (report: InitProgressReport) => void; /** - * Cache to store model related data. + * Cache to store model related data, implemented with the Cache API. */ export class ArtifactCache implements ArtifactCacheTemplate { private scope: string; @@ -1008,7 +1008,7 @@ export class ArtifactCache implements ArtifactCacheTemplate { /** * Convert the Response object to the expected storetype instead - * @param response the cache or indexDB response + * @param response the cache or indexedDB response * @param storetype the storetype stored in the indexedDB database * @returns the expected response object */ @@ -1028,7 +1028,7 @@ export class ArtifactCache implements ArtifactCacheTemplate { /** * fetch the corresponding url object in response or stored object format * @param url url - * @param storetype the storage type for indexDB + * @param storetype the storage type for indexedDB * @returns response in json, arraybuffer or pure response format */ async fetchWithCache(url: string, storetype?: string) { @@ -1094,9 +1094,9 @@ export class ArtifactCache implements ArtifactCacheTemplate { } /** - * Cache by IndexDB to support caching model data + * Cache by IndexedDB to support caching model data */ -export class ArtifactindexDBCache implements ArtifactCacheTemplate { +export class ArtifactIndexedDBCache implements ArtifactCacheTemplate { private dbName?: string; private dbVersion = 1; private db: IDBDatabase | undefined; @@ -1191,7 +1191,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } } - async addToIndexDB(url: string, response: any, storetype?: string) { + async addToIndexedDB(url: string, response: any, storetype?: string) { await this.initDB(); let data: any; if (storetype != undefined) { @@ -1200,7 +1200,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { } else if (storetype.toLocaleLowerCase() === "arraybuffer") { data = await response.arrayBuffer(); } else { - console.error("Unsupported Type in IndexDB"); + console.error("Unsupported Type in IndexedDB"); } } return new Promise((resolve, reject) => { @@ -1223,7 +1223,7 @@ export class ArtifactindexDBCache implements ArtifactCacheTemplate { throw new Error('Network response was not ok'); } const response_copy = response.clone(); - await this.addToIndexDB(url, response_copy, storetype); + await this.addToIndexedDB(url, response_copy, storetype); if (storetype.toLowerCase() === "arraybuffer") { return await response.arrayBuffer(); } else if (storetype.toLowerCase() === "json") { @@ -1733,7 +1733,7 @@ export class Instance implements Disposable { * @param ndarrayCacheUrl The cache url. * @param device The device to be fetched to. * @param cacheScope The scope identifier of the cache - * @param cacheType The type of the cache: "cache" or "indexDB" + * @param cacheType The type of the cache: "cache" or "indexedDB" * @returns The meta data */ async fetchNDArrayCache( @@ -1748,8 +1748,8 @@ export class Instance implements Disposable { } if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexdb") { - artifactCache = new ArtifactindexDBCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexeddb") { + artifactCache = new ArtifactIndexedDBCache(cacheScope); } else { console.error("Unsupported Cache Type, using default browser cache"); artifactCache = new ArtifactCache(cacheScope); @@ -2838,7 +2838,7 @@ export function instantiate( * * @param ndarrayCacheUrl The cache url which links to the NDArray * @param cacheScope The scope identifier of the cache - * @param cacheType The type of the cache: "cache" or "indexDB" + * @param cacheType The type of the cache: "cache" or "indexedDB" * @returns the result if the cache has NDArray */ export async function hasNDArrayInCache( @@ -2849,8 +2849,8 @@ export async function hasNDArrayInCache( let artifactCache; if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexdb") { - artifactCache = new ArtifactindexDBCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexeddb") { + artifactCache = new ArtifactIndexedDBCache(cacheScope); } else { console.error("Unsupported Cache Type, using default browser cache"); artifactCache = new ArtifactCache(cacheScope); @@ -2877,7 +2877,7 @@ export async function hasNDArrayInCache( * * @param cacheUrl The cacheUrl for the items * @param cacheScope The scope identifier of the cache - * @param cacheType The type of the cache: "cache" or "indexDB" + * @param cacheType The type of the cache: "cache" or "indexedDB" */ export async function deleteNDArrayCache( cacheUrl: string, @@ -2887,8 +2887,8 @@ export async function deleteNDArrayCache( let artifactCache; if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexdb") { - artifactCache = new ArtifactindexDBCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexeddb") { + artifactCache = new ArtifactIndexedDBCache(cacheScope); } else { console.error("Unsupported Cache Type, using default browser cache"); artifactCache = new ArtifactCache(cacheScope); From 47310cc16cc3e2c629b3bf2019e622589f172298 Mon Sep 17 00:00:00 2001 From: Charlie Ruan <53290280+CharlieFRuan@users.noreply.github.com> Date: Fri, 5 Apr 2024 23:21:40 -0400 Subject: [PATCH 4/6] Modify addToCache and fetchWithCache logics - We make addToCache not return anything, see new specification - This allows us to skip downloaded files in fetchNDArrayCache instead of running into DOMException: Key already exists - Call addToCache in fetchWithCache, so we only need to retrieve afterwards - Remove cacheOnly for callback, use loading instead (since we separated download and loading) - Fix responseTostoretype bug in ArtifactCache --- web/src/artifact_cache.ts | 43 +++++++---- web/src/runtime.ts | 158 +++++++++++++------------------------- 2 files changed, 82 insertions(+), 119 deletions(-) diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts index f1bf55921941..fc44c7ea2873 100644 --- a/web/src/artifact_cache.ts +++ b/web/src/artifact_cache.ts @@ -21,32 +21,45 @@ */ export interface ArtifactCacheTemplate { /** - * fetch key url from cache, optional storetype for IndexedDB + * Retrieve data object that corresponds to `url` from cache. If data object does not exist in + * cache, fetch the data and then add to cache. * - * storagetype for indexedDB have two options: - * @param url: return a json object - * 2. arraybuffer: return an arraybuffer object + * @param url: The url to the data to be cached. + * @param storetype: This field is required so that `ArtifactIndexedDBCache` can store the + * actual data object (see `addToCache()`), while `ArtifactCache` which uses the Cache API can + * return the actual data object rather than the request. There are two options: + * 1. "json": returns equivalent to `fetch(url).json()` + * 2. "arraybuffer": returns equivalent to `fetch(url).arraybuffer()` + * @return The data object (i.e. users do not need to call `.json()` or `.arraybuffer()`). + * + * @note This is an async function. */ - fetchWithCache(url: string, storetype?: string); + fetchWithCache(url: string, storetype?: string): Promise; /** - * add key url to cache, optional storetype for IndexedDB + * Fetch data from url and add into cache. If already exists in cache, should return instantly. * - * storagetype for indexedDB have two options: - * 1. json: return a json object - * 2. arraybuffer: return an arraybuffer object - * - * returns the response or the specified stored object - * for reduced database transaction + * @param url: The url to the data to be cached. + * @param storetype: Only applies to `ArtifactIndexedDBCache`. Since `indexedDB` stores the actual + * data rather than a request, we specify `storagetype`. There are two options: + * 1. "json": IndexedDB stores `fetch(url).json()` + * 2. "arraybuffer": IndexedDB stores `fetch(url).arrayBuffer()` + * + * @note This is an async function. */ - addToCache(url: string, storetype?: string): Promise; + addToCache(url: string, storetype?: string): Promise; + /** * check if cache has all keys in Cache + * + * @note This is an async function. */ - hasAllKeys(keys: string[]); + hasAllKeys(keys: string[]): Promise; /** * Delete url in cache if url exists + * + * @note This is an async function. */ - deleteInCache(url: string); + deleteInCache(url: string): Promise; } diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 0d15520dfce3..8280066a8424 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -989,7 +989,6 @@ export interface NDArrayShardEntry { export interface InitProgressReport { progress: number; timeElapsed: number; - cacheOnly: boolean; text: string; } @@ -1008,21 +1007,18 @@ export class ArtifactCache implements ArtifactCacheTemplate { /** * Convert the Response object to the expected storetype instead - * @param response the cache or indexedDB response - * @param storetype the storetype stored in the indexedDB database - * @returns the expected response object - */ - async responseTostoretype(response: Response, storetype: string): Promise { - let result: any; - if (storetype.toLowerCase() === "json") { - result = await result.json(); + */ + async responseTostoretype(response: Response, storetype?: string): Promise { + if (storetype === undefined) { + return response; + } else if (storetype.toLowerCase() === "json") { + return await response.json(); } else if (storetype.toLowerCase() === "arraybuffer") { - result = await result.arrayBuffer(); + return await response.arrayBuffer(); } else { - console.error("Unknown storage type, return raw response"); + console.error("Unknown storage type " + storetype + ", returning raw response"); return response; } - return result; } /** @@ -1031,38 +1027,25 @@ export class ArtifactCache implements ArtifactCacheTemplate { * @param storetype the storage type for indexedDB * @returns response in json, arraybuffer or pure response format */ - async fetchWithCache(url: string, storetype?: string) { - const request = new Request(url); - if (this.cache === undefined) { - this.cache = await caches.open(this.scope); - } - let result = await this.cache.match(request); + async fetchWithCache(url: string, storetype?: string): Promise { + await this.addToCache(url, storetype); + const result = await this.cache.match(new Request(url)); if (result === undefined) { - result = await this.addToCache(url, storetype); - return result; - } else { - if (storetype === undefined) { - return result; - } else { - return await this.responseTostoretype(result, storetype); - } + // Already called `addToCache()`, should expect the request in cache. + throw Error("Cannot fetch " + url); } + return await this.responseTostoretype(result, storetype); } + // eslint-disable-next-line @typescript-eslint/no-unused-vars async addToCache(url: string, storetype?: string) { const request = new Request(url); if (this.cache === undefined) { this.cache = await caches.open(this.scope); } - let result = await this.cache.match(request); + const result = await this.cache.match(request); if (result === undefined) { await this.cache.add(request); - result = await this.cache.match(request); - } - if (storetype === undefined) { - return result; - } else { - return await this.responseTostoretype(result, storetype); } } @@ -1153,7 +1136,7 @@ export class ArtifactIndexedDBCache implements ArtifactCacheTemplate { }); } - async asyncGetHelper(url: string) { + async asyncGetHelper(url: string): Promise { return new Promise((resolve, reject) => { let result: any; const transaction = this.db?.transaction(['urls'], 'readonly'); @@ -1170,37 +1153,33 @@ export class ArtifactIndexedDBCache implements ArtifactCacheTemplate { }) } - async fetchWithCache(url: string, storetype?: string) { - await this.initDB(); // await the initDB process - const isInDB = await this.isUrlInDB(url); - if (!isInDB) { - const response = await this.addToCache(url, storetype); - return response; - } else { - // URL found in DB, just fetch without storing - const result = await this.asyncGetHelper(url); - if (result != null && typeof result === "object" && "data" in result) { - return result.data; - } else if (result === null) { - // previously null data in cache! - await this.deleteInCache(url); - const response = await this.addToCache(url, storetype); - return response; - } - return null; + async fetchWithCache(url: string, storetype?: string): Promise { + await this.addToCache(url, storetype); + let result = await this.asyncGetHelper(url); + if (result === null) { + // previously null data in cache or somehow failed to add to cache, delete and retry + await this.deleteInCache(url); + await this.addToCache(url, storetype); + result = await this.asyncGetHelper(url); } + if (result != null && typeof result === "object" && "data" in result) { + // `storetype` not used here because the data stored in indexedDB is already in that type + return result.data; + } + throw Error("ArtifactIndexedDBCache failed to fetch: " + url); } async addToIndexedDB(url: string, response: any, storetype?: string) { await this.initDB(); let data: any; + // IndexedDB, unlike the Cache API, stores the actual data object, so we convert reponse here. if (storetype != undefined) { if (storetype.toLowerCase() === "json") { data = await response.json(); } else if (storetype.toLocaleLowerCase() === "arraybuffer") { data = await response.arrayBuffer(); } else { - console.error("Unsupported Type in IndexedDB"); + throw Error("Unsupported storetyp for IndexedDB: " + storetype); } } return new Promise((resolve, reject) => { @@ -1215,24 +1194,22 @@ export class ArtifactIndexedDBCache implements ArtifactCacheTemplate { }); } - async addToCache(url: string, storetype?: string): Promise { - let response: Response; + async addToCache(url: string, storetype?: string): Promise { + await this.initDB(); // await the initDB process + // If already cached, nothing to do + const isInDB = await this.isUrlInDB(url); + if (isInDB) { + return; + } try { - response = await fetch(url); + const response = await fetch(url); if (!response.ok) { throw new Error('Network response was not ok'); } const response_copy = response.clone(); await this.addToIndexedDB(url, response_copy, storetype); - if (storetype.toLowerCase() === "arraybuffer") { - return await response.arrayBuffer(); - } else if (storetype.toLowerCase() === "json") { - return await response.json(); - } else { - return response; - } } catch (error) { - console.error("There was a problem fetching the data:", error); + throw Error("Failed to store " + url + " with error: " + error); } } @@ -1742,26 +1719,17 @@ export class Instance implements Disposable { cacheScope = "tvmjs", cacheType = "cache" ): Promise { - let artifactCache; - if (cacheType === undefined) { - artifactCache = new ArtifactCache(cacheScope); - } - if (cacheType.toLowerCase() === "cache") { + let artifactCache: ArtifactCacheTemplate; + if (cacheType === undefined || cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); } else if (cacheType.toLowerCase() == "indexeddb") { artifactCache = new ArtifactIndexedDBCache(cacheScope); } else { - console.error("Unsupported Cache Type, using default browser cache"); + console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); artifactCache = new ArtifactCache(cacheScope); } const jsonUrl = new URL("ndarray-cache.json", ndarrayCacheUrl).href; - const result = await artifactCache.fetchWithCache(jsonUrl, "json"); - let list; - if (result instanceof Response) { - list = await result.json(); - } else { - list = result; - } + const list = await artifactCache.fetchWithCache(jsonUrl, "json"); await this.fetchNDArrayCacheInternal( ndarrayCacheUrl, list["records"] as Array, device, artifactCache); @@ -1794,13 +1762,13 @@ export class Instance implements Disposable { let timeElapsed = 0; const cacheOnly = await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)); + + // `loading`: we have finished downloading (or already cacheOnly) and are loading onto WebGPU const reportCallback = (iter: number, loading = false) => { // report for (let j = 0; j < this.initProgressCallback.length; ++j) { let text: string; if (loading) { - text = "Finished fetching params, loading onto WebGPU."; - } else 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, " @@ -1816,7 +1784,6 @@ export class Instance implements Disposable { this.initProgressCallback[j]({ progress: fetchedBytes / totalBytes, timeElapsed: timeElapsed, - cacheOnly: cacheOnly, text: text }); } @@ -1826,7 +1793,6 @@ export class Instance implements Disposable { this.initProgressCallback[j]({ progress: fetchedBytes / totalBytes, timeElapsed: 0, - cacheOnly: cacheOnly, text: "Start to fetch params", }); } @@ -1845,7 +1811,7 @@ export class Instance implements Disposable { } timeElapsed = Math.ceil((perf.now() - tstart) / 1000); fetchedBytes += shard.nbytes; - reportCallback(fetchedShards++); + reportCallback(fetchedShards++, /*loading=*/false); } } // We launch 4 parallel for loops to limit the max concurrency to 4 download @@ -1857,7 +1823,6 @@ export class Instance implements Disposable { downloadCache(2 * loopSize, 3 * loopSize), downloadCache(3 * loopSize, list.length) ]); - reportCallback(list.length, /*loading=*/true); } // Then iteratively, load the shard from cache @@ -1908,9 +1873,7 @@ export class Instance implements Disposable { throw err; } } - if (cacheOnly) { - reportCallback(i + 1, /* Need to Report call back*/false); - } + reportCallback(i + 1, /*loading=*/true); } } @@ -2383,7 +2346,6 @@ export class Instance implements Disposable { this.initProgressCallback[j]({ progress: progress, timeElapsed: timeElapsed, - cacheOnly: false, text: text }); } @@ -2846,13 +2808,13 @@ export async function hasNDArrayInCache( cacheScope = "tvmjs", cacheType = "cache" ): Promise { - let artifactCache; + let artifactCache: ArtifactCacheTemplate; if (cacheType.toLowerCase() === "cache") { artifactCache = new ArtifactCache(cacheScope); } else if (cacheType.toLowerCase() == "indexeddb") { artifactCache = new ArtifactIndexedDBCache(cacheScope); } else { - console.error("Unsupported Cache Type, using default browser cache"); + console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); artifactCache = new ArtifactCache(cacheScope); } const jsonUrl = new URL("ndarray-cache.json", ndarrayCacheUrl).href; @@ -2860,13 +2822,7 @@ export async function hasNDArrayInCache( if (!hasJsonUrlInCache) { return false; } - const result = await artifactCache.fetchWithCache(jsonUrl, "json"); - let list; - if (result instanceof Response) { - list = await result.json(); - } else { - list = result; - } + let list = await artifactCache.fetchWithCache(jsonUrl, "json"); list = list["records"] as Array; return await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)); } @@ -2890,17 +2846,11 @@ export async function deleteNDArrayCache( } else if (cacheType.toLowerCase() == "indexeddb") { artifactCache = new ArtifactIndexedDBCache(cacheScope); } else { - console.error("Unsupported Cache Type, using default browser cache"); + console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); artifactCache = new ArtifactCache(cacheScope); } const jsonUrl = new URL("ndarray-cache.json", cacheUrl).href; - const result = await artifactCache.fetchWithCache(jsonUrl, "json"); - let list; - if (result instanceof Response) { - list = await result.json(); - } else { - list = result; - } + const list = await artifactCache.fetchWithCache(jsonUrl, "json"); const arrayentry = list["records"] as Array; const processShard = async (i: number) => { const dataUrl = new URL(arrayentry[i].dataPath, cacheUrl).href; From b9b06fc00eb85d556f61c19d7f180e166b5db9d2 Mon Sep 17 00:00:00 2001 From: Charlie Ruan <53290280+CharlieFRuan@users.noreply.github.com> Date: Fri, 5 Apr 2024 23:41:56 -0400 Subject: [PATCH 5/6] Move all cache related code to artifact_cache --- web/src/artifact_cache.ts | 351 +++++++++++++++++++++++++++++++++++++ web/src/index.ts | 9 +- web/src/runtime.ts | 355 +------------------------------------- 3 files changed, 365 insertions(+), 350 deletions(-) diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts index fc44c7ea2873..cd965358d2be 100644 --- a/web/src/artifact_cache.ts +++ b/web/src/artifact_cache.ts @@ -16,6 +16,23 @@ * specific language governing permissions and limitations * under the License. */ + +export interface NDArrayCacheEntry { + name: string; + shape: Array; + dtype: string; + format: "f32-to-bf16" | "raw"; + byteOffset: number; + nbytes: number; +} + +export interface NDArrayShardEntry { + dataPath: string; + format: "raw-shard"; + nbytes: number; + records: Array; +} + /** * Common Interface for the artifact cache */ @@ -63,3 +80,337 @@ export interface ArtifactCacheTemplate { */ deleteInCache(url: string): Promise; } + + +/** + * Cache to store model related data, implemented with the Cache API. + */ +export class ArtifactCache implements ArtifactCacheTemplate { + private scope: string; + private cache?: Cache; + + constructor(scope: string) { + this.scope = scope; + } + + /** + * Convert the Response object to the expected storetype instead + */ + async responseTostoretype(response: Response, storetype?: string): Promise { + if (storetype === undefined) { + return response; + } else if (storetype.toLowerCase() === "json") { + return await response.json(); + } else if (storetype.toLowerCase() === "arraybuffer") { + return await response.arrayBuffer(); + } else { + console.error("Unknown storage type " + storetype + ", returning raw response"); + return response; + } + } + + /** + * fetch the corresponding url object in response or stored object format + * @param url url + * @param storetype the storage type for indexedDB + * @returns response in json, arraybuffer or pure response format + */ + async fetchWithCache(url: string, storetype?: string): Promise { + await this.addToCache(url, storetype); + const result = await this.cache.match(new Request(url)); + if (result === undefined) { + // Already called `addToCache()`, should expect the request in cache. + throw Error("Cannot fetch " + url); + } + return await this.responseTostoretype(result, storetype); + } + + // eslint-disable-next-line @typescript-eslint/no-unused-vars + async addToCache(url: string, storetype?: string) { + const request = new Request(url); + if (this.cache === undefined) { + this.cache = await caches.open(this.scope); + } + const result = await this.cache.match(request); + if (result === undefined) { + await this.cache.add(request); + } + } + + /** + * Determine if all keys exist in the cache + * @param keys the url key list of the strings + * @returns boolean value indicate if all keys are in cache + */ + 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(() => false); + } + + /** + * Delete the corresponding url object in cache + * @param url the corresponding url object to be deleted + */ + async deleteInCache(url: string) { + if (this.cache === undefined) { + this.cache = await caches.open(this.scope); + } + await this.cache.delete(url); + } +} + +/** + * Cache by IndexedDB to support caching model data + */ +export class ArtifactIndexedDBCache implements ArtifactCacheTemplate { + private dbName?: string; + private dbVersion = 1; + private db: IDBDatabase | undefined; + + constructor(dbName: string) { + this.dbName = dbName; + } + + /** + * Init the indexed DB database if it is not initialized. + */ + private async initDB() { + if (this.db != null) { + return; // the db is already inialized + } + return new Promise((resolve, reject) => { + const request = indexedDB.open(this.dbName, this.dbVersion); + request.onupgradeneeded = (event) => { + this.db = (event.target as IDBOpenDBRequest).result; + if (!this.db.objectStoreNames.contains('urls')) { + this.db.createObjectStore('urls', { keyPath: 'url' }); + } + }; + request.onsuccess = (event) => { + this.db = (event.target as IDBOpenDBRequest).result; + resolve(); + }; + request.onerror = (event) => { + console.error("Database error: ", (event.target as IDBOpenDBRequest).error); + reject((event.target as IDBOpenDBRequest).error); + }; + }); + } + + /** + * Check if current url object is in indexedDB or not + * @param url the url link + * @returns boolean indicate if url object in indexedDB + */ + private async isUrlInDB(url: string): Promise { + return new Promise((resolve, reject) => { + const transaction = this.db?.transaction(['urls'], 'readonly'); + if (transaction === undefined) { + return false; + } + const store = transaction.objectStore('urls'); + const request = store.get(url); + request.onsuccess = () => { + resolve(request.result !== undefined); + }; + request.onerror = (event) => { + reject((event.target as IDBRequest).error); + }; + }); + } + + async asyncGetHelper(url: string): Promise { + return new Promise((resolve, reject) => { + let result: any; + const transaction = this.db?.transaction(['urls'], 'readonly'); + if (transaction === undefined) { + return false; + } + transaction.oncomplete = () => resolve(result); + transaction.onerror = () => reject(transaction.error); + const objectStore = transaction.objectStore('urls'); + const getRequest = objectStore.get(url); + getRequest.onsuccess = () => { + result = getRequest.result; + } + }) + } + + async fetchWithCache(url: string, storetype?: string): Promise { + await this.addToCache(url, storetype); + let result = await this.asyncGetHelper(url); + if (result === null) { + // previously null data in cache or somehow failed to add to cache, delete and retry + await this.deleteInCache(url); + await this.addToCache(url, storetype); + result = await this.asyncGetHelper(url); + } + if (result != null && typeof result === "object" && "data" in result) { + // `storetype` not used here because the data stored in indexedDB is already in that type + return result.data; + } + throw Error("ArtifactIndexedDBCache failed to fetch: " + url); + } + + async addToIndexedDB(url: string, response: any, storetype?: string) { + await this.initDB(); + let data: any; + // IndexedDB, unlike the Cache API, stores the actual data object, so we convert reponse here. + if (storetype != undefined) { + if (storetype.toLowerCase() === "json") { + data = await response.json(); + } else if (storetype.toLocaleLowerCase() === "arraybuffer") { + data = await response.arrayBuffer(); + } else { + throw Error("Unsupported storetyp for IndexedDB: " + storetype); + } + } + return new Promise((resolve, reject) => { + const transaction = this.db?.transaction(['urls'], 'readwrite'); + if (transaction === undefined) { + return; + } + const store = transaction.objectStore('urls'); + const request = store.add({ data, url }); // Index DB follows a {value, key} format, instead of {key, value} format! + request.onsuccess = () => resolve(); + request.onerror = (event) => reject((event.target as IDBRequest).error); + }); + } + + async addToCache(url: string, storetype?: string): Promise { + await this.initDB(); // await the initDB process + // If already cached, nothing to do + const isInDB = await this.isUrlInDB(url); + if (isInDB) { + return; + } + try { + const response = await fetch(url); + if (!response.ok) { + throw new Error('Network response was not ok'); + } + const response_copy = response.clone(); + await this.addToIndexedDB(url, response_copy, storetype); + } catch (error) { + throw Error("Failed to store " + url + " with error: " + error); + } + } + + async hasAllKeys(keys: string[]): Promise { + await this.initDB(); // Ensure the DB is initialized + if (!this.db) { + throw new Error('Database is not initialized'); + } + return new Promise((resolve, reject) => { + const transaction = this.db.transaction(['urls'], 'readonly'); + const store = transaction.objectStore('urls'); + const promises = keys.map(key => { + return new Promise((resolve) => { + const request = store.get(key); + request.onsuccess = () => { + if (request.result === undefined) { + resolve(false); // Key not found, resolve with false + } else { + resolve(true); // Key found, resolve with true + } + }; + request.onerror = () => { + resolve(false); // On error, resolve as if the key was not found + }; + }); + }); + Promise.all(promises).then(results => { + const allExist = results.every(exists => exists); + resolve(allExist); + }).catch(error => { + reject(error); // Reject the main promise if any of the promises are rejected + }); + }); + } + + async deleteInCache(url: string) { + await this.initDB(); // Make sure the DB is initialized + const transaction = this.db?.transaction(['urls'], 'readwrite'); + if (transaction === undefined) { + return; + } + const store = transaction.objectStore('urls'); + const request = store.delete(url); + // Await completion of the delete request + await new Promise((resolve, reject) => { + request.onsuccess = () => resolve(); + request.onerror = () => reject(request.error); + }); + return; + } +} + + +/** + * Function to check if NDarray is in Cache or not + * + * @param ndarrayCacheUrl The cache url which links to the NDArray + * @param cacheScope The scope identifier of the cache + * @param cacheType The type of the cache: "cache" or "indexedDB" + * @returns the result if the cache has NDArray + */ +export async function hasNDArrayInCache( + ndarrayCacheUrl: string, + cacheScope = "tvmjs", + cacheType = "cache" +): Promise { + let artifactCache: ArtifactCacheTemplate; + if (cacheType.toLowerCase() === "cache") { + artifactCache = new ArtifactCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexeddb") { + artifactCache = new ArtifactIndexedDBCache(cacheScope); + } else { + console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); + artifactCache = new ArtifactCache(cacheScope); + } + const jsonUrl = new URL("ndarray-cache.json", ndarrayCacheUrl).href; + const hasJsonUrlInCache = await artifactCache.hasAllKeys([jsonUrl]); + if (!hasJsonUrlInCache) { + return false; + } + let list = await artifactCache.fetchWithCache(jsonUrl, "json"); + list = list["records"] as Array; + return await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)); +} + + +/** + * Given cacheUrl, search up items to delete based on cacheUrl/ndarray-cache.json + * + * @param cacheUrl The cacheUrl for the items + * @param cacheScope The scope identifier of the cache + * @param cacheType The type of the cache: "cache" or "indexedDB" + */ +export async function deleteNDArrayCache( + cacheUrl: string, + cacheScope = "tvmjs", + cacheType = "cache" +) { + let artifactCache: ArtifactCacheTemplate; + if (cacheType.toLowerCase() === "cache") { + artifactCache = new ArtifactCache(cacheScope); + } else if (cacheType.toLowerCase() == "indexeddb") { + artifactCache = new ArtifactIndexedDBCache(cacheScope); + } else { + console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); + artifactCache = new ArtifactCache(cacheScope); + } + const jsonUrl = new URL("ndarray-cache.json", cacheUrl).href; + const list = await artifactCache.fetchWithCache(jsonUrl, "json"); + const arrayentry = list["records"] as Array; + const processShard = async (i: number) => { + const dataUrl = new URL(arrayentry[i].dataPath, cacheUrl).href; + await artifactCache.deleteInCache(dataUrl); + } + await Promise.all(arrayentry.map((_, index) => processShard(index))); +} diff --git a/web/src/index.ts b/web/src/index.ts index f134b0b5dc0b..d4fc9b9187e6 100644 --- a/web/src/index.ts +++ b/web/src/index.ts @@ -22,8 +22,15 @@ export { PackedFunc, Module, NDArray, TVMArray, TVMObject, VirtualMachine, InitProgressCallback, InitProgressReport, - ArtifactCache, ArtifactIndexedDBCache, Instance, instantiate, hasNDArrayInCache, deleteNDArrayCache + Instance, instantiate } from "./runtime"; +export { + ArtifactCacheTemplate, + ArtifactCache, + ArtifactIndexedDBCache, + hasNDArrayInCache, + deleteNDArrayCache +} from "./artifact_cache"; export { Disposable, LibraryProvider } from "./types"; export { RPCServer } from "./rpc_server"; export { assert, wasmPath, LinearCongruentialGenerator } from "./support"; diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 8280066a8424..4b40bbc34152 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -27,8 +27,12 @@ import { assert, StringToUint8Array, LinearCongruentialGenerator } from "./suppo import { Environment } from "./environment"; import { AsyncifyHandler } from "./asyncify"; import { FunctionInfo, WebGPUContext } from "./webgpu"; -import { ArtifactCacheTemplate } from "./artifact_cache"; - +import { + ArtifactCache, + ArtifactCacheTemplate, + ArtifactIndexedDBCache, + NDArrayShardEntry, +} from "./artifact_cache"; import * as compact from "./compact"; import * as ctypes from "./ctypes"; @@ -970,21 +974,6 @@ enum AsyncCallbackCode { kReturn = 4, kException = 5, } -export interface NDArrayCacheEntry { - name: string; - shape: Array; - dtype: string; - format: "f32-to-bf16" | "raw"; - byteOffset: number; - nbytes: number; -} - -export interface NDArrayShardEntry { - dataPath: string; - format: "raw-shard"; - nbytes: number; - records: Array; -} export interface InitProgressReport { progress: number; @@ -994,274 +983,6 @@ export interface InitProgressReport { export type InitProgressCallback = (report: InitProgressReport) => void; -/** - * Cache to store model related data, implemented with the Cache API. - */ -export class ArtifactCache implements ArtifactCacheTemplate { - private scope: string; - private cache?: Cache; - - constructor(scope: string) { - this.scope = scope; - } - - /** - * Convert the Response object to the expected storetype instead - */ - async responseTostoretype(response: Response, storetype?: string): Promise { - if (storetype === undefined) { - return response; - } else if (storetype.toLowerCase() === "json") { - return await response.json(); - } else if (storetype.toLowerCase() === "arraybuffer") { - return await response.arrayBuffer(); - } else { - console.error("Unknown storage type " + storetype + ", returning raw response"); - return response; - } - } - - /** - * fetch the corresponding url object in response or stored object format - * @param url url - * @param storetype the storage type for indexedDB - * @returns response in json, arraybuffer or pure response format - */ - async fetchWithCache(url: string, storetype?: string): Promise { - await this.addToCache(url, storetype); - const result = await this.cache.match(new Request(url)); - if (result === undefined) { - // Already called `addToCache()`, should expect the request in cache. - throw Error("Cannot fetch " + url); - } - return await this.responseTostoretype(result, storetype); - } - - // eslint-disable-next-line @typescript-eslint/no-unused-vars - async addToCache(url: string, storetype?: string) { - const request = new Request(url); - if (this.cache === undefined) { - this.cache = await caches.open(this.scope); - } - const result = await this.cache.match(request); - if (result === undefined) { - await this.cache.add(request); - } - } - - /** - * Determine if all keys exist in the cache - * @param keys the url key list of the strings - * @returns boolean value indicate if all keys are in cache - */ - 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(() => false); - } - - /** - * Delete the corresponding url object in cache - * @param url the corresponding url object to be deleted - */ - async deleteInCache(url: string) { - if (this.cache === undefined) { - this.cache = await caches.open(this.scope); - } - await this.cache.delete(url); - } -} - -/** - * Cache by IndexedDB to support caching model data - */ -export class ArtifactIndexedDBCache implements ArtifactCacheTemplate { - private dbName?: string; - private dbVersion = 1; - private db: IDBDatabase | undefined; - - constructor(dbName: string) { - this.dbName = dbName; - } - - /** - * Init the indexed DB database if it is not initialized. - */ - private async initDB() { - if (this.db != null) { - return; // the db is already inialized - } - return new Promise((resolve, reject) => { - const request = indexedDB.open(this.dbName, this.dbVersion); - request.onupgradeneeded = (event) => { - this.db = (event.target as IDBOpenDBRequest).result; - if (!this.db.objectStoreNames.contains('urls')) { - this.db.createObjectStore('urls', { keyPath: 'url' }); - } - }; - request.onsuccess = (event) => { - this.db = (event.target as IDBOpenDBRequest).result; - resolve(); - }; - request.onerror = (event) => { - console.error("Database error: ", (event.target as IDBOpenDBRequest).error); - reject((event.target as IDBOpenDBRequest).error); - }; - }); - } - - /** - * Check if current url object is in indexedDB or not - * @param url the url link - * @returns boolean indicate if url object in indexedDB - */ - private async isUrlInDB(url: string): Promise { - return new Promise((resolve, reject) => { - const transaction = this.db?.transaction(['urls'], 'readonly'); - if (transaction === undefined) { - return false; - } - const store = transaction.objectStore('urls'); - const request = store.get(url); - request.onsuccess = () => { - resolve(request.result !== undefined); - }; - request.onerror = (event) => { - reject((event.target as IDBRequest).error); - }; - }); - } - - async asyncGetHelper(url: string): Promise { - return new Promise((resolve, reject) => { - let result: any; - const transaction = this.db?.transaction(['urls'], 'readonly'); - if (transaction === undefined) { - return false; - } - transaction.oncomplete = () => resolve(result); - transaction.onerror = () => reject(transaction.error); - const objectStore = transaction.objectStore('urls'); - const getRequest = objectStore.get(url); - getRequest.onsuccess = () => { - result = getRequest.result; - } - }) - } - - async fetchWithCache(url: string, storetype?: string): Promise { - await this.addToCache(url, storetype); - let result = await this.asyncGetHelper(url); - if (result === null) { - // previously null data in cache or somehow failed to add to cache, delete and retry - await this.deleteInCache(url); - await this.addToCache(url, storetype); - result = await this.asyncGetHelper(url); - } - if (result != null && typeof result === "object" && "data" in result) { - // `storetype` not used here because the data stored in indexedDB is already in that type - return result.data; - } - throw Error("ArtifactIndexedDBCache failed to fetch: " + url); - } - - async addToIndexedDB(url: string, response: any, storetype?: string) { - await this.initDB(); - let data: any; - // IndexedDB, unlike the Cache API, stores the actual data object, so we convert reponse here. - if (storetype != undefined) { - if (storetype.toLowerCase() === "json") { - data = await response.json(); - } else if (storetype.toLocaleLowerCase() === "arraybuffer") { - data = await response.arrayBuffer(); - } else { - throw Error("Unsupported storetyp for IndexedDB: " + storetype); - } - } - return new Promise((resolve, reject) => { - const transaction = this.db?.transaction(['urls'], 'readwrite'); - if (transaction === undefined) { - return; - } - const store = transaction.objectStore('urls'); - const request = store.add({ data, url }); // Index DB follows a {value, key} format, instead of {key, value} format! - request.onsuccess = () => resolve(); - request.onerror = (event) => reject((event.target as IDBRequest).error); - }); - } - - async addToCache(url: string, storetype?: string): Promise { - await this.initDB(); // await the initDB process - // If already cached, nothing to do - const isInDB = await this.isUrlInDB(url); - if (isInDB) { - return; - } - try { - const response = await fetch(url); - if (!response.ok) { - throw new Error('Network response was not ok'); - } - const response_copy = response.clone(); - await this.addToIndexedDB(url, response_copy, storetype); - } catch (error) { - throw Error("Failed to store " + url + " with error: " + error); - } - } - - async hasAllKeys(keys: string[]): Promise { - await this.initDB(); // Ensure the DB is initialized - if (!this.db) { - throw new Error('Database is not initialized'); - } - return new Promise((resolve, reject) => { - const transaction = this.db.transaction(['urls'], 'readonly'); - const store = transaction.objectStore('urls'); - const promises = keys.map(key => { - return new Promise((resolve) => { - const request = store.get(key); - request.onsuccess = () => { - if (request.result === undefined) { - resolve(false); // Key not found, resolve with false - } else { - resolve(true); // Key found, resolve with true - } - }; - request.onerror = () => { - resolve(false); // On error, resolve as if the key was not found - }; - }); - }); - Promise.all(promises).then(results => { - const allExist = results.every(exists => exists); - resolve(allExist); - }).catch(error => { - reject(error); // Reject the main promise if any of the promises are rejected - }); - }); - } - - async deleteInCache(url: string) { - await this.initDB(); // Make sure the DB is initialized - const transaction = this.db?.transaction(['urls'], 'readwrite'); - if (transaction === undefined) { - return; - } - const store = transaction.objectStore('urls'); - const request = store.delete(url); - // Await completion of the delete request - await new Promise((resolve, reject) => { - request.onsuccess = () => resolve(); - request.onerror = () => reject(request.error); - }); - return; - } -} - /** * TVM runtime instance. * @@ -2794,67 +2515,3 @@ export function instantiate( } ); } - -/** - * Function to check if NDarray is in Cache or not - * - * @param ndarrayCacheUrl The cache url which links to the NDArray - * @param cacheScope The scope identifier of the cache - * @param cacheType The type of the cache: "cache" or "indexedDB" - * @returns the result if the cache has NDArray - */ -export async function hasNDArrayInCache( - ndarrayCacheUrl: string, - cacheScope = "tvmjs", - cacheType = "cache" -): Promise { - let artifactCache: ArtifactCacheTemplate; - if (cacheType.toLowerCase() === "cache") { - artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexeddb") { - artifactCache = new ArtifactIndexedDBCache(cacheScope); - } else { - console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); - artifactCache = new ArtifactCache(cacheScope); - } - const jsonUrl = new URL("ndarray-cache.json", ndarrayCacheUrl).href; - const hasJsonUrlInCache = await artifactCache.hasAllKeys([jsonUrl]); - if (!hasJsonUrlInCache) { - return false; - } - let list = await artifactCache.fetchWithCache(jsonUrl, "json"); - list = list["records"] as Array; - return await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, ndarrayCacheUrl).href)); -} - - -/** - * Given cacheUrl, search up items to delete based on cacheUrl/ndarray-cache.json - * - * @param cacheUrl The cacheUrl for the items - * @param cacheScope The scope identifier of the cache - * @param cacheType The type of the cache: "cache" or "indexedDB" - */ -export async function deleteNDArrayCache( - cacheUrl: string, - cacheScope = "tvmjs", - cacheType = "cache" -) { - let artifactCache; - if (cacheType.toLowerCase() === "cache") { - artifactCache = new ArtifactCache(cacheScope); - } else if (cacheType.toLowerCase() == "indexeddb") { - artifactCache = new ArtifactIndexedDBCache(cacheScope); - } else { - console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache."); - artifactCache = new ArtifactCache(cacheScope); - } - const jsonUrl = new URL("ndarray-cache.json", cacheUrl).href; - const list = await artifactCache.fetchWithCache(jsonUrl, "json"); - const arrayentry = list["records"] as Array; - const processShard = async (i: number) => { - const dataUrl = new URL(arrayentry[i].dataPath, cacheUrl).href; - await artifactCache.deleteInCache(dataUrl); - } - await Promise.all(arrayentry.map((_, index) => processShard(index))); -} From 081094b3fe6a9692b9630fbe21aecbf3d19ebfc9 Mon Sep 17 00:00:00 2001 From: Charlie Ruan <53290280+CharlieFRuan@users.noreply.github.com> Date: Sat, 6 Apr 2024 00:04:51 -0400 Subject: [PATCH 6/6] Fix lint --- web/src/artifact_cache.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts index cd965358d2be..f833df1be523 100644 --- a/web/src/artifact_cache.ts +++ b/web/src/artifact_cache.ts @@ -48,7 +48,7 @@ export interface ArtifactCacheTemplate { * 1. "json": returns equivalent to `fetch(url).json()` * 2. "arraybuffer": returns equivalent to `fetch(url).arraybuffer()` * @return The data object (i.e. users do not need to call `.json()` or `.arraybuffer()`). - * + * * @note This is an async function. */ fetchWithCache(url: string, storetype?: string): Promise; @@ -61,21 +61,21 @@ export interface ArtifactCacheTemplate { * data rather than a request, we specify `storagetype`. There are two options: * 1. "json": IndexedDB stores `fetch(url).json()` * 2. "arraybuffer": IndexedDB stores `fetch(url).arrayBuffer()` - * + * * @note This is an async function. */ addToCache(url: string, storetype?: string): Promise; /** * check if cache has all keys in Cache - * + * * @note This is an async function. */ hasAllKeys(keys: string[]): Promise; /** * Delete url in cache if url exists - * + * * @note This is an async function. */ deleteInCache(url: string): Promise;