[Web] Let external owners defer device-loss disposal

Signed-off-by: Akaash Parthasarathy <akaashrp@gmail.com>

Author: akaashrp

web/src/runtime.ts

@@ -921,6 +921,7 @@ export class Instance implements Disposable {
   private initProgressCallback: Array<InitProgressCallback> = [];
   private rng: LinearCongruentialGenerator;
   private deviceLostIsError = true;  // whether device.lost is due to actual error or dispose()
+  private autoDisposeOnDeviceLost = true;
   private cacheState: CacheState = new CacheState();
 
   /**
@@ -2080,7 +2081,7 @@ export class Instance implements Disposable {
     });
 
     device.lost.then((info: any) => {
-      if (this.deviceLostIsError) {
+      if (this.deviceLostIsError && this.autoDisposeOnDeviceLost) {
         console.error("Device lost, calling Instance.dispose(). Please initialize again. ", info);
         this.dispose();
       }
@@ -2108,6 +2109,17 @@ export class Instance implements Disposable {
     this.lib.webGPUContext = webGPUContext;
   }
 
+  /**
+   * Configure automatic disposal after WebGPU device loss.
+   *
+   * External owners should disable automatic disposal after initialization if
+   * they serialize disposal with active runtime calls.
+   * @param enabled Whether device loss should immediately dispose this instance.
+   */
+  setDeviceLostAutoDispose(enabled: boolean): void {
+    this.autoDisposeOnDeviceLost = enabled;
+  }
+
   /** Register all object factory */
   private registerObjectFactoryFuncs(): void {
     this.registerObjectConstructor("ffi.Array",

web/tests/node/test_tensor_cache_webgpu.js

@@ -39,7 +39,18 @@ function createInstance() {
   );
 }
 
-function createMockGPUDevice({ detachWriteSources = false } = {}) {
+function createDeferred() {
+  let resolve;
+  const promise = new Promise((resolvePromise) => {
+    resolve = resolvePromise;
+  });
+  return { promise, resolve };
+}
+
+function createMockGPUDevice({
+  detachWriteSources = false,
+  lost = new Promise(() => {}),
+} = {}) {
   const buffers = [];
   const writes = [];
   const queue = {
@@ -65,7 +76,7 @@ function createMockGPUDevice({ detachWriteSources = false } = {}) {
   };
   const device = {
     queue,
-    lost: new Promise(() => {}),
+    lost,
     addEventListener: jest.fn(),
     pushErrorScope: jest.fn(),
     popErrorScope: jest.fn(() => Promise.resolve(null)),
@@ -94,6 +105,33 @@ function createArtifactCache(manifest, shard) {
   };
 }
 
+test("an external owner can defer disposal after device loss", async () => {
+  const lost = createDeferred();
+  const tvm = createInstance();
+  const dispose = jest.spyOn(tvm, "dispose");
+  const gpu = createMockGPUDevice({ lost: lost.promise });
+  const log = jest.spyOn(console, "error").mockImplementation(() => {});
+  try {
+    tvm.initWebGPU(gpu.device);
+    tvm.setDeviceLostAutoDispose(false);
+
+    lost.resolve({ reason: "unknown", message: "test device loss" });
+    await lost.promise;
+    await Promise.resolve();
+
+    expect(dispose).not.toHaveBeenCalled();
+
+    tvm.dispose();
+    expect(dispose).toHaveBeenCalledTimes(1);
+  } finally {
+    if (dispose.mock.calls.length === 0) {
+      tvm.dispose();
+    }
+    dispose.mockRestore();
+    log.mockRestore();
+  }
+});
+
 test("WebGPU tensor cache uploads pass-through records and decodes BF16 in place", async () => {
   const tvm = createInstance();
   const gpu = createMockGPUDevice();