From 5772a00020d0af7574117433eca80593539caaab Mon Sep 17 00:00:00 2001 From: Nestor Qin Date: Fri, 31 May 2024 20:15:20 -0400 Subject: [PATCH] feat: fall back to Web Worker if WebGPU is not available in Service Worker --- app/client/webllm.ts | 13 +++++---- app/components/home.tsx | 56 +++++++++++++++++++++++++----------- app/worker/service-worker.ts | 35 ++++++++++++++++++++-- package.json | 1 + tsconfig.json | 2 +- yarn.lock | 5 ++++ 6 files changed, 86 insertions(+), 26 deletions(-) diff --git a/app/client/webllm.ts b/app/client/webllm.ts index 6fe71cbd..7607efaa 100644 --- a/app/client/webllm.ts +++ b/app/client/webllm.ts @@ -34,17 +34,18 @@ export class WebLLMApi implements LLMApi { private initialized = false; webllm: WebLLMHandler; - constructor(logLevel: LogLevel = "WARN") { - if ("serviceWorker" in navigator && navigator.serviceWorker.controller) { - log.info("Service Worker API is available and in use."); + constructor( + type: "serviceWorker" | "webWorker", + logLevel: LogLevel = "WARN", + ) { + if (type === "serviceWorker") { + log.info("Create ServiceWorkerMLCEngine"); this.webllm = { type: "serviceWorker", engine: new ServiceWorkerMLCEngine(KEEP_ALIVE_INTERVAL), }; } else { - log.info( - "Service Worker API is unavailable. Falling back to use web worker.", - ); + log.info("Create WebWorkerMLCEngine"); this.webllm = { type: "webWorker", engine: new WebWorkerMLCEngine( diff --git a/app/components/home.tsx b/app/components/home.tsx index 33ee60b6..c67c4d81 100644 --- a/app/components/home.tsx +++ b/app/components/home.tsx @@ -167,36 +167,60 @@ const useWebLLM = () => { // Initialize WebLLM engine useEffect(() => { if ("serviceWorker" in navigator) { + log.info("Service Worker API is available and in use."); navigator.serviceWorker.ready.then(() => { - setWebLLM(new WebLLMApi(config.logLevel)); - - // Verify service worker has been activated - const heartbeatCallback = (event: MessageEvent) => { - const msg = event.data; - if (msg.kind === "heartbeat") { - log.info("Confirmed messages from Service Worker. Starting app..."); - setWebllmAlive(true); + log.info("Service Worker is activated."); + // Check whether WebGPU is available in Service Worker + const request = { + kind: "checkWebGPUAvilability", + uuid: crypto.randomUUID(), + content: "", + }; + const webGPUCheckCallback = (event: MessageEvent) => { + const message = event.data; + if (message.kind === "return" && message.uuid === request.uuid) { + const isWebGPUAvailable = message.content; + log.info( + isWebGPUAvailable + ? "Service Worker has WebGPU Available." + : "Service Worker does not have available WebGPU.", + ); + if (!webllm && !isWebllmActive) { + setWebLLM( + new WebLLMApi( + isWebGPUAvailable ? "serviceWorker" : "webWorker", + config.logLevel, + ), + ); + setWebllmAlive(true); + } navigator.serviceWorker.removeEventListener( "message", - heartbeatCallback, + webGPUCheckCallback, ); } }; - navigator.serviceWorker.addEventListener("message", heartbeatCallback); - navigator.serviceWorker.controller?.postMessage({ - kind: "keepAlive", - uuid: crypto.randomUUID(), - }); + navigator.serviceWorker.addEventListener( + "message", + webGPUCheckCallback, + ); + navigator.serviceWorker.controller?.postMessage(request); }); } else { - setWebLLM(new WebLLMApi(config.logLevel)); + log.info( + "Service Worker API is unavailable. Falling back to use web worker.", + ); + setWebLLM(new WebLLMApi("webWorker", config.logLevel)); setWebllmAlive(true); } // If service worker registration timeout setTimeout(() => { if (!navigator.serviceWorker?.controller && !isWebllmActive && !webllm) { - setWebLLM(new WebLLMApi(config.logLevel)); + log.info( + "Service Worker activation is timed out. Falling back to use web worker.", + ); + setWebLLM(new WebLLMApi("webWorker", config.logLevel)); setWebllmAlive(true); } }, 5_000); diff --git a/app/worker/service-worker.ts b/app/worker/service-worker.ts index ceb19826..b3446898 100644 --- a/app/worker/service-worker.ts +++ b/app/worker/service-worker.ts @@ -2,17 +2,46 @@ import { ServiceWorkerMLCEngineHandler, MLCEngine } from "@neet-nestor/web-llm"; import { defaultCache } from "@serwist/next/worker"; import type { PrecacheEntry, SerwistGlobalConfig } from "serwist"; import { CacheFirst, ExpirationPlugin, Serwist } from "serwist"; -import log from "loglevel"; declare const self: ServiceWorkerGlobalScope; const CHATGPT_NEXT_WEB_CACHE = "chatgpt-next-web-cache"; const engine = new MLCEngine(); let handler: ServiceWorkerMLCEngineHandler; +async function checkGPUAvailablity() { + if (!("gpu" in navigator)) { + console.log("Service Worker: Web-LLM Engine Activated"); + return false; + } + const adapter = await navigator.gpu.requestAdapter(); + if (!adapter) { + console.log("Service Worker: Web-LLM Engine Activated"); + return false; + } + return true; +} + self.addEventListener("message", (event) => { if (!handler) { handler = new ServiceWorkerMLCEngineHandler(engine); - log.info("Service Worker: Web-LLM Engine Activated"); + console.log("Service Worker: Web-LLM Engine Activated"); + } + + const msg = event.data; + if (msg.kind === "checkWebGPUAvilability") { + console.log("Service Worker: Web-LLM Engine Activated"); + checkGPUAvailablity().then((gpuAvailable) => { + console.log( + "Service Worker: WebGPU is " + + (gpuAvailable ? "available" : "unavailable"), + ); + const reply = { + kind: "return", + uuid: msg.uuid, + content: gpuAvailable, + }; + event.source?.postMessage(reply); + }); } }); @@ -30,7 +59,7 @@ self.addEventListener("install", (event) => { self.addEventListener("activate", (event) => { if (!handler) { handler = new ServiceWorkerMLCEngineHandler(engine); - log.info("Service Worker: Web-LLM Engine Activated"); + console.log("Service Worker: Web-LLM Engine Activated"); } }); diff --git a/package.json b/package.json index 84246d6e..92f65435 100644 --- a/package.json +++ b/package.json @@ -49,6 +49,7 @@ "@types/react-dom": "^18.2.7", "@types/react-katex": "^3.0.0", "@types/spark-md5": "^3.0.4", + "@webgpu/types": "^0.1.42", "cross-env": "^7.0.3", "eslint": "^8.49.0", "eslint-config-next": "13.4.19", diff --git a/tsconfig.json b/tsconfig.json index ca215e30..ed47bdde 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -1,7 +1,7 @@ { "compilerOptions": { "target": "ES2015", - "types": ["@serwist/next/typings"], + "types": ["@serwist/next/typings", "@webgpu/types"], "lib": ["dom", "dom.iterable", "esnext", "webworker"], "allowJs": true, "skipLibCheck": true, diff --git a/yarn.lock b/yarn.lock index 68ed03c3..80fb650d 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2889,6 +2889,11 @@ "@webassemblyjs/ast" "1.11.6" "@xtuc/long" "4.2.2" +"@webgpu/types@^0.1.42": + version "0.1.42" + resolved "https://registry.yarnpkg.com/@webgpu/types/-/types-0.1.42.tgz#02ce82ae7d486521064eda2e9c7e38697136afe9" + integrity sha512-uvJtt4OD1Vjdebrrz3kNLgpOicYbikwnM8WPG6YD2lkCOHDtPdEtCINJFIFtbOCtPfA8SreR/vKyUNbAt92IwQ== + "@xtuc/ieee754@^1.2.0": version "1.2.0" resolved "https://registry.npmmirror.com/@xtuc/ieee754/-/ieee754-1.2.0.tgz#eef014a3145ae477a1cbc00cd1e552336dceb790"