WaveCut/Bonsai-Chat-WebGPU

1import { describe, expect, it, vi } from 'vitest';2import { NativeLogCollector } from './native-log';3import { BrowserEngineRuntime } from './runtime';4 5interface RuntimeInternals {6  wllama: { exit: () => Promise<void> } | null;7  loaded: {8    backend: 'webgpu' | 'wasm';9    model: { id: string; cpuFallback: boolean };10  } | null;11  nativeLog: NativeLogCollector;12  raceWebGpuDeviceLoss<Result>(13    stage: 'load' | 'generate' | 'backend-report',14    backend: 'webgpu' | 'wasm',15    model: { id: string; cpuFallback: boolean },16    operation: Promise<Result>,17  ): Promise<Result>;18}19 20describe('BrowserEngineRuntime WebGPU device loss', () => {21  it('invalidates the model and Wllama instance before returning a reload steer', async () => {22    const runtime = new BrowserEngineRuntime();23    const internals = runtime as unknown as RuntimeInternals;24    const exit = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);25    internals.wllama = { exit };26    internals.loaded = {27      backend: 'webgpu',28      model: { id: '1_7b', cpuFallback: true },29    };30    internals.nativeLog.append([31      'ggml_webgpu: Device lost! Reason: 2, Message: GPU process crashed',32    ]);33 34    await expect(runtime.backendReport()).rejects.toMatchObject({35      name: 'EngineRuntimeError',36      code: 'WEBGPU_DEVICE_LOST',37      details: {38        recoverable: true,39        nextAction: 'reload-model',40        stage: 'backend-report',41        modelId: '1_7b',42        cpuFallbackAvailable: true,43        cause: 'device-loss',44        signal: { reason: 2, message: 'GPU process crashed' },45      },46    });47    expect(exit).toHaveBeenCalledOnce();48    expect(internals.loaded).toBeNull();49    expect(internals.wllama).toBeNull();50  });51 52  it('labels an out-of-memory device loss and tells the user to reduce context', async () => {53    const runtime = new BrowserEngineRuntime();54    const internals = runtime as unknown as RuntimeInternals;55    internals.wllama = { exit: vi.fn<() => Promise<void>>().mockResolvedValue(undefined) };56    internals.loaded = {57      backend: 'webgpu',58      model: { id: '27b', cpuFallback: false },59    };60    internals.nativeLog.append([61      'ggml_webgpu: Device lost! Reason: 2, Message: Out of device memory',62    ]);63 64    await expect(runtime.backendReport()).rejects.toMatchObject({65      code: 'WEBGPU_DEVICE_LOST',66      message: expect.stringContaining('Reload with a smaller context'),67      details: { cause: 'memory-pressure', stage: 'backend-report' },68    });69  });70 71  it('does not invalidate a CPU-WASM model for a stale WebGPU log signal', async () => {72    const runtime = new BrowserEngineRuntime();73    const internals = runtime as unknown as RuntimeInternals;74    const exit = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);75    internals.wllama = { exit };76    internals.loaded = {77      backend: 'wasm',78      model: { id: '1_7b', cpuFallback: true },79    };80    internals.nativeLog.append([81      'ggml_webgpu: Device lost! Reason: 1, Message: unknown',82    ]);83 84    await expect(runtime.backendReport()).resolves.toEqual({85      backends: [],86      nGraphSplits: null,87      opsOnCpu: 0,88      layersGpu: null,89      flashAttention: null,90      cacheTypeK: null,91      cacheTypeV: null,92      webgpuKvBufferBytes: null,93      webgpuTrace: [],94    });95    expect(exit).not.toHaveBeenCalled();96    expect(internals.loaded).not.toBeNull();97  });98 99  it('interrupts a hung native operation as soon as the log callback signals device loss', async () => {100    const runtime = new BrowserEngineRuntime();101    const internals = runtime as unknown as RuntimeInternals;102    const exit = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);103    const model = { id: '1_7b', cpuFallback: true };104    internals.wllama = { exit };105    internals.loaded = { backend: 'webgpu', model };106    const hung = new Promise<never>(() => undefined);107    const operation = internals.raceWebGpuDeviceLoss('generate', 'webgpu', model, hung);108 109    internals.nativeLog.append([110      'ggml_webgpu: Device lost! Reason: 4, Message: device removed',111    ]);112 113    await expect(operation).rejects.toMatchObject({114      code: 'WEBGPU_DEVICE_LOST',115      details: { stage: 'generate', nextAction: 'reload-model' },116    });117    expect(exit).toHaveBeenCalledOnce();118    expect(internals.loaded).toBeNull();119    expect(internals.wllama).toBeNull();120  });121});122