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