Migrate cross origin storage implementation to tvm/web

Author: akaashrp

README.md

@@ -150,6 +150,29 @@ const engine = new MLCEngine({
 await engine.reload(selectedModel);
 ```
 
+### Cache Backend Policy
+
+WebLLM supports three cache backends through `AppConfig.cacheBackend`:
+
+- `"cache"`: browser [Cache API](https://developer.mozilla.org/en-US/docs/Web/API/Cache) (default).
+- `"indexeddb"`: browser [IndexedDB](https://developer.mozilla.org/en-US/docs/Web/API/IndexedDB_API).
+- `"cross-origin"`: experimental Chrome [Cross-Origin Storage](https://github.com/explainers-by-googlers/cross-origin-storage) extension backend.
+
+Example:
+
+```typescript
+import { CreateMLCEngine, prebuiltAppConfig } from "@mlc-ai/web-llm";
+
+const appConfig = { ...prebuiltAppConfig, cacheBackend: "cross-origin" };
+const engine = await CreateMLCEngine("Llama-3.1-8B-Instruct-q4f32_1-MLC", {
+  appConfig,
+});
+```
+
+Notes:
+- The `"cross-origin"` backend requires installing and enabling a compatible browser extension.
+- Cross-origin backend currently does not support programmatic tensor-cache deletion; clearing is extension-managed.
+
 ### Chat Completion
 After successfully initializing the engine, you can now invoke chat completions using OpenAI style chat APIs through the `engine.chat.completions` interface. For the full list of parameters and their descriptions, check [section below](#full-openai-compatibility) and [OpenAI API reference](https://platform.openai.com/docs/api-reference/chat/create).
 

examples/README.md

@@ -47,7 +47,7 @@ These examples demonstrate various capabilities via WebLLM's OpenAI-like API.
 
 - [logit-processor](logit-processor): while `logit_bias` is supported, we additionally support stateful logit processing where users can specify their own rules. We also expose low-level API `forwardTokensAndSample()`.
 - [cache-usage](cache-usage): demonstrates how WebLLM supports multiple cache backends. Choose between the [Cache API](https://developer.mozilla.org/en-US/docs/Web/API/Cache), [IndexedDB cache](https://developer.mozilla.org/en-US/docs/Web/API/IndexedDB_API), or the experimental Chrome [Cross-Origin Storage](https://github.com/explainers-by-googlers/cross-origin-storage) extension via `appConfig.cacheBackend`. Also demonstrates various cache utils such as checking
-  whether a model is cached, deleting a model's weights from cache, deleting a model library wasm from cache, etc.
+  whether a model is cached, deleting a model's weights from cache, deleting a model library wasm from cache, etc. Note: cross-origin backend currently does not support programmatic tensor-cache deletion.
 - [simple-chat-upload](simple-chat-upload): demonstrates how to upload local models to WebLLM instead of downloading via a URL link
 
 ## Demo Spaces

examples/cache-usage/README.md

@@ -5,9 +5,10 @@ setting `AppConfig.cacheBackend` to `"cache"`, `"indexeddb"`, or `"cross-origin"
 This folder provides an example on how different caches are used in WebLLM. We also
 demonstrate the utility cache functions such as deleting models, checking if models are in cache, etc.
 
-> **Note:** The cross-origin backend requires installation of the [cross-origin storage browser extension](https://github.com/web-ai-community/cross-origin-storage-extension).
+> **Note:** The cross-origin backend requires installation of the [cross-origin storage browser extension](https://github.com/web-ai-community/cross-origin-storage-extension). This does not currently support programmatic tensor-cache deletion; deletion is extension-managed.
 
-For more information about the two caches, see: https://developer.mozilla.org/en-US/docs/Web/API/Storage_API/Storage_quotas_and_eviction_criteria#what_technologies_store_data_in_the_browser.
+For more information about Cache API and IndexedDB, see:
+https://developer.mozilla.org/en-US/docs/Web/API/Storage_API/Storage_quotas_and_eviction_criteria#what_technologies_store_data_in_the_browser.
 
 To inspect the downloaded artifacts in your browser, open up developer console, go to application,
 and you will find the artifacts under either `IndexedDB` or `Cache storage`. When `"cross-origin"` is selected,

examples/simple-chat-upload/src/simple_chat.ts

@@ -395,7 +395,7 @@ async function uploadFiles(): Promise<void> {
     alert("No files selected.");
     return;
   }
-  if (appConfig.useIndexedDBCache) {
+  if (appConfig.cacheBackend === "indexeddb") {
     for (const file of input.files) {
       uploadToIndexedDB(file);
     }

src/cache_util.ts

@@ -9,151 +9,27 @@ import {
 import { cleanModelUrl } from "./support";
 import { ModelNotFoundError, UnsupportedTokenizerFilesError } from "./error";
 import { Tokenizer } from "@mlc-ai/web-tokenizers";
-import CrossOriginStorage from "./cross_origin_storage";
-import CrossOriginStorageCache from "./cross_origin_storage_cache";
 
 type CacheScope = "webllm/model" | "webllm/config" | "webllm/wasm";
 
-let crossOriginUnavailableLogged = false;
-let crossOriginAvailabilityWait: Promise<void> | null = null;
-
-function scheduleCrossOriginFallbackWarning(
-  logger: (msg: string) => void,
-): void {
-  if (crossOriginUnavailableLogged || crossOriginAvailabilityWait) {
-    return;
-  }
-  crossOriginAvailabilityWait = (async () => {
-    const available = CrossOriginStorage.isAvailable();
-    crossOriginAvailabilityWait = null;
-    if (available || crossOriginUnavailableLogged) {
-      return;
-    }
-    logger(
-      "Cross-origin storage backend is not yet available; temporarily falling back to the Cache API.",
-    );
-    crossOriginUnavailableLogged = true;
-  })();
-}
-
-function useCrossOrigin(appConfig: AppConfig): boolean {
-  return (
-    getCacheBackend(appConfig) === "cross-origin" &&
-    CrossOriginStorage.isAvailable()
-  );
-}
-
-export function getArtifactCache(
+function getCacheAccessOptions(
   scope: CacheScope,
   appConfig: AppConfig,
-  logger: (msg: string) => void = console.warn,
-): tvmjs.ArtifactCacheTemplate {
-  const backend = getCacheBackend(appConfig);
-  if (backend === "cross-origin") {
-    if (CrossOriginStorage.isAvailable()) {
-      return new CrossOriginStorageCache(scope);
-    }
-    scheduleCrossOriginFallbackWarning(logger);
-  }
-  if (backend === "indexeddb") {
-    return new tvmjs.ArtifactIndexedDBCache(scope);
-  }
-  return new tvmjs.ArtifactCache(scope);
-}
-
-async function hasTensorCache(
-  cache: tvmjs.ArtifactCacheTemplate,
-  tensorCacheUrl: string,
-): Promise<boolean> {
-  const jsonUrl = new URL("tensor-cache.json", tensorCacheUrl).href;
-  const hasManifest = await cache.hasAllKeys([jsonUrl]);
-  if (!hasManifest) {
-    return false;
-  }
-  const manifest = await cache.fetchWithCache(jsonUrl, "json");
-  const records = manifest?.records ?? [];
-  if (!Array.isArray(records) || records.length === 0) {
-    return false;
-  }
-  const shardUrls = records.map(
-    (entry: { dataPath: string }) =>
-      new URL(entry.dataPath, tensorCacheUrl).href,
-  );
-  return cache.hasAllKeys(shardUrls);
+): tvmjs.TensorCacheAccessOptions {
+  return {
+    cacheScope: scope,
+    cacheType: getCacheBackend(appConfig),
+  };
 }
 
-async function deleteTensorCacheEntries(
-  cache: tvmjs.ArtifactCacheTemplate,
-  tensorCacheUrl: string,
-): Promise<void> {
-  const jsonUrl = new URL("tensor-cache.json", tensorCacheUrl).href;
-  const hasManifest = await cache.hasAllKeys([jsonUrl]);
-  if (!hasManifest) {
-    return;
-  }
-  let manifest: { records?: Array<{ dataPath: string }> };
-  try {
-    manifest = await cache.fetchWithCache(jsonUrl, "json");
-  } catch (err) {
-    console.warn(
-      `Failed to load tensor cache manifest at ${jsonUrl}; skipping deletion.`,
-      err,
-    );
-    return;
-  }
-  const records = manifest?.records ?? [];
-  await Promise.all(
-    records.map(async (entry) => {
-      if (!entry?.dataPath) {
-        return;
-      }
-      const dataUrl = new URL(entry.dataPath, tensorCacheUrl).href;
-      await cache.deleteInCache(dataUrl);
-    }),
-  );
-  await cache.deleteInCache(jsonUrl);
-}
-
-export async function fetchModelArtifacts(
-  tvm: tvmjs.Instance,
-  tensorCacheUrl: string,
-  device: tvmjs.DLDevice,
+function createScopedArtifactCache(
+  scope: CacheScope,
   appConfig: AppConfig,
-  signal?: AbortSignal,
-): Promise<any> {
-  if (!useCrossOrigin(appConfig)) {
-    const backend = getCacheBackend(appConfig);
-    const cacheType = backend === "indexeddb" ? "indexeddb" : "cache";
-    return tvm.fetchTensorCache(
-      tensorCacheUrl,
-      device,
-      "webllm/model",
-      cacheType,
-      signal,
-    );
-  }
-
-  const artifactCache = getArtifactCache("webllm/model", appConfig);
-  const jsonUrl = new URL("tensor-cache.json", tensorCacheUrl).href;
-  const manifest = await artifactCache.fetchWithCache(jsonUrl, "json", signal);
-  const records = (
-    Array.isArray(manifest?.records) ? manifest.records : []
-  ) as Array<any>;
-  await (tvm as any).fetchTensorCacheInternal(
-    tensorCacheUrl,
-    records,
-    device,
-    artifactCache,
-    signal,
+): tvmjs.ArtifactCacheTemplate {
+  return tvmjs.createArtifactCache(
+    scope,
+    getCacheAccessOptions(scope, appConfig),
   );
-  if (manifest?.metadata !== undefined) {
-    const runtime = tvm as any;
-    runtime.cacheMetadata = {
-      ...runtime.cacheMetadata,
-      ...(manifest.metadata as Record<string, unknown>),
-    };
-  }
-  return manifest;
 }
 
 function findModelRecord(modelId: string, appConfig?: AppConfig): ModelRecord {
@@ -175,13 +51,10 @@ export async function hasModelInCache(
   }
   const modelRecord = findModelRecord(modelId, appConfig);
   const modelUrl = cleanModelUrl(modelRecord.model);
-  if (useCrossOrigin(appConfig)) {
-    const cache = getArtifactCache("webllm/model", appConfig);
-    return hasTensorCache(cache, modelUrl);
-  }
-  const backend = getCacheBackend(appConfig);
-  const cacheType = backend === "indexeddb" ? "indexeddb" : "cache";
-  return tvmjs.hasTensorInCache(modelUrl, "webllm/model", cacheType);
+  return tvmjs.hasTensorInCache(
+    modelUrl,
+    getCacheAccessOptions("webllm/model", appConfig),
+  );
 }
 
 export async function deleteModelAllInfoInCache(
@@ -210,14 +83,11 @@ export async function deleteModelInCache(
   }
   const modelRecord = findModelRecord(modelId, appConfig);
   const modelUrl = cleanModelUrl(modelRecord.model);
-  const modelCache = getArtifactCache("webllm/model", appConfig);
-  if (useCrossOrigin(appConfig)) {
-    await deleteTensorCacheEntries(modelCache, modelUrl);
-  } else {
-    const backend = getCacheBackend(appConfig);
-    const cacheType = backend === "indexeddb" ? "indexeddb" : "cache";
-    await tvmjs.deleteTensorCache(modelUrl, "webllm/model", cacheType);
-  }
+  const modelCache = createScopedArtifactCache("webllm/model", appConfig);
+  await tvmjs.deleteTensorCache(
+    modelUrl,
+    getCacheAccessOptions("webllm/model", appConfig),
+  );
   await modelCache.deleteInCache(new URL("tokenizer.model", modelUrl).href);
   await modelCache.deleteInCache(new URL("tokenizer.json", modelUrl).href);
 }
@@ -231,7 +101,7 @@ export async function deleteChatConfigInCache(
     appConfig = prebuiltAppConfig;
   }
   const modelRecord = findModelRecord(modelId, appConfig);
-  const configCache = getArtifactCache("webllm/config", appConfig);
+  const configCache = createScopedArtifactCache("webllm/config", appConfig);
   const modelUrl = cleanModelUrl(modelRecord.model);
   const configUrl = new URL("mlc-chat-config.json", modelUrl).href;
   await configCache.deleteInCache(configUrl);
@@ -246,7 +116,7 @@ export async function deleteModelWasmInCache(
     appConfig = prebuiltAppConfig;
   }
   const modelRecord = findModelRecord(modelId, appConfig);
-  const wasmCache = getArtifactCache("webllm/wasm", appConfig);
+  const wasmCache = createScopedArtifactCache("webllm/wasm", appConfig);
   await wasmCache.deleteInCache(modelRecord.model_lib);
 }
 
@@ -264,7 +134,7 @@ export async function asyncLoadTokenizer(
   appConfig: AppConfig,
   logger: (msg: string) => void = console.log,
 ): Promise<Tokenizer> {
-  const modelCache = getArtifactCache("webllm/model", appConfig, logger);
+  const modelCache = createScopedArtifactCache("webllm/model", appConfig);
 
   if (config.tokenizer_files.includes("tokenizer.json")) {
     const url = new URL("tokenizer.json", baseUrl).href;

src/config.ts

@@ -272,6 +272,10 @@ export interface ModelRecord {
  * @param cacheBackend: the backend to use for caching models and other artifacts.
  * If unspecified, will use the Cache API. For more information, see:
  * https://developer.mozilla.org/en-US/docs/Web/API/Storage_API/Storage_quotas_and_eviction_criteria#what_technologies_store_data_in_the_browser
+ * Supported values are:
+ * - "cache": browser Cache API.
+ * - "indexeddb": IndexedDB-backed cache.
+ * - "cross-origin": Chrome Cross-Origin Storage extension-backed cache.
  *
  * @note Note that the Cache API is more well-tested in WebLLM as of now.
  */