feat(stage-ui): should handle WebGPU device-loss (#1680)
--------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -51,7 +51,7 @@
|
||||
},
|
||||
"inlinedDependencies": {
|
||||
"@electron-toolkit/preload": "3.0.2",
|
||||
"@moeru/eventa": "catalog:",
|
||||
"@moeru/eventa": "1.0.0-beta.4",
|
||||
"async-mutex": "0.5.0",
|
||||
"nanoid": [
|
||||
"5.1.6",
|
||||
|
||||
@@ -15,10 +15,13 @@ vi.mock('../../../composables/use-inference-status', () => ({
|
||||
removeInferenceStatus: vi.fn(),
|
||||
}))
|
||||
|
||||
const recordDeviceLoss = vi.fn()
|
||||
vi.mock('../coordinator', () => ({
|
||||
getGPUCoordinator: () => ({
|
||||
requestAllocation: vi.fn(() => ({ modelId: 'test', estimatedBytes: 0 })),
|
||||
release: vi.fn(),
|
||||
touch: vi.fn(),
|
||||
recordDeviceLoss,
|
||||
}),
|
||||
getLoadQueue: () => ({
|
||||
enqueue: vi.fn((_id: string, _p: number, loader: () => Promise<unknown>) => loader()),
|
||||
@@ -76,3 +79,33 @@ describe('classifyError phase integration', () => {
|
||||
expect(classifyError(new Error('tensor shape mismatch'), 'inference')).toBe('INFERENCE_FAILED')
|
||||
})
|
||||
})
|
||||
|
||||
describe('kokoro adapter - device loss resilience', () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers()
|
||||
recordDeviceLoss.mockClear()
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers()
|
||||
vi.restoreAllMocks()
|
||||
})
|
||||
|
||||
it('should start with zero device-loss count and null manifest', async () => {
|
||||
const { createKokoroAdapter } = await import('./kokoro')
|
||||
const adapter = createKokoroAdapter()
|
||||
|
||||
expect(adapter.deviceLossCount).toBe(0)
|
||||
expect(adapter.manifest).toBeNull()
|
||||
})
|
||||
|
||||
it('should expose manifest and deviceLossCount as readonly getters', async () => {
|
||||
const { createKokoroAdapter } = await import('./kokoro')
|
||||
const adapter = createKokoroAdapter()
|
||||
|
||||
// Getters should be defined and callable
|
||||
expect(typeof adapter.deviceLossCount).toBe('number')
|
||||
// Manifest is explicitly null before any load
|
||||
expect(adapter.manifest).toBeNull()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -13,10 +13,10 @@ import { defaultPerfTracer } from '@proj-airi/stage-shared'
|
||||
import { Mutex } from 'async-mutex'
|
||||
|
||||
import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status'
|
||||
import { MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants'
|
||||
import { DEVICE_LOSS_WASM_THRESHOLD, MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants'
|
||||
import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator'
|
||||
import { LOAD_PRIORITY } from '../load-queue'
|
||||
import { classifyError, createRequestId } from '../protocol'
|
||||
import { classifyDeviceLossReason, classifyError, createRequestId } from '../protocol'
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Types
|
||||
@@ -41,6 +41,16 @@ export interface KokoroAdapter {
|
||||
|
||||
/** Current state */
|
||||
readonly state: 'idle' | 'loading' | 'ready' | 'running' | 'error' | 'terminated'
|
||||
|
||||
/**
|
||||
* Snapshot of the last successful load config, or null if never loaded.
|
||||
* `device` reflects the device actually used (post WASM promotion / worker
|
||||
* fallback), which may differ from the device requested by the caller.
|
||||
*/
|
||||
readonly manifest: { quantization: string, device: string } | null
|
||||
|
||||
/** Number of WebGPU device-loss events observed by this adapter */
|
||||
readonly deviceLossCount: number
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -153,6 +163,11 @@ function waitForWorkerMessage<T = any>(
|
||||
// Factory
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
interface KokoroManifest {
|
||||
quantization: string
|
||||
device: string
|
||||
}
|
||||
|
||||
export function createKokoroAdapter(): KokoroAdapter {
|
||||
let worker: Worker | null = null
|
||||
let state: KokoroAdapter['state'] = 'idle'
|
||||
@@ -162,6 +177,13 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
let currentModelStatusId: string | null = null
|
||||
let errorListener: ((event: ErrorEvent) => void) | null = null
|
||||
|
||||
// NOTICE: Device-loss resilience state. `lastManifest` records the last
|
||||
// successful load config so scheduleRestart can reconstruct context if the
|
||||
// worker died. `deviceLossCount` tracks WebGPU device-loss events so we
|
||||
// can promote to WASM after repeated failures (see DEVICE_LOSS_WASM_THRESHOLD).
|
||||
let lastManifest: KokoroManifest | null = null
|
||||
let deviceLossCount = 0
|
||||
|
||||
const operationMutex = new Mutex()
|
||||
const lifecycleMutex = new Mutex()
|
||||
|
||||
@@ -174,9 +196,22 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
worker.addEventListener('error', errorListener)
|
||||
}
|
||||
|
||||
function handleWorkerError(_event: ErrorEvent | Error): void {
|
||||
function handleWorkerError(event: ErrorEvent | Error): void {
|
||||
state = 'error'
|
||||
operationMutex.cancel()
|
||||
|
||||
// Record device-loss telemetry before teardown so the coordinator sees it
|
||||
// even if the adapter is never used again.
|
||||
const code = classifyError(event instanceof Error ? event : (event as ErrorEvent).error ?? event)
|
||||
if (code === 'DEVICE_LOST') {
|
||||
deviceLossCount++
|
||||
getGPUCoordinator().recordDeviceLoss({
|
||||
modelId: currentModelStatusId ?? MODEL_NAMES.KOKORO,
|
||||
reason: classifyDeviceLossReason(event instanceof Error ? event : (event as ErrorEvent).error ?? event),
|
||||
occurredAt: Date.now(),
|
||||
})
|
||||
}
|
||||
|
||||
destroyWorker()
|
||||
scheduleRestart()
|
||||
}
|
||||
@@ -237,6 +272,22 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
device: string,
|
||||
options?: { onProgress?: (p: ProgressPayload) => void },
|
||||
): Promise<Voices> {
|
||||
// NOTICE: Proactive WASM promotion. If this adapter has suffered repeated
|
||||
// WebGPU device-loss events, webgpu is unreliable on this device and we
|
||||
// should not keep retrying. The worker's per-load dtype/device fallback
|
||||
// chain handles transient failures; this guard handles persistent ones.
|
||||
let effectiveDevice = device
|
||||
if (
|
||||
device === 'webgpu'
|
||||
&& deviceLossCount >= DEVICE_LOSS_WASM_THRESHOLD
|
||||
) {
|
||||
console.warn(
|
||||
`[KokoroAdapter] ${deviceLossCount} device-loss events recorded, `
|
||||
+ `promoting load from webgpu to wasm.`,
|
||||
)
|
||||
effectiveDevice = 'wasm'
|
||||
}
|
||||
|
||||
await ensureStarted()
|
||||
|
||||
return defaultPerfTracer.withMeasure('inference', 'kokoro-load-model', () => operationMutex.runExclusive(async () => {
|
||||
@@ -248,7 +299,7 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
removeInferenceStatus(currentModelStatusId)
|
||||
currentModelStatusId = modelStatusId
|
||||
|
||||
updateInferenceStatus(modelStatusId, { state: 'downloading', device: device as any })
|
||||
updateInferenceStatus(modelStatusId, { state: 'downloading', device: effectiveDevice as any })
|
||||
|
||||
// Use the global load queue to serialize model loads across all adapters
|
||||
return getLoadQueue().enqueue(modelStatusId, LOAD_PRIORITY.TTS, async () => {
|
||||
@@ -275,7 +326,7 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
type: 'load-model',
|
||||
requestId,
|
||||
modelId: MODEL_NAMES.KOKORO,
|
||||
device,
|
||||
device: effectiveDevice,
|
||||
dtype: quantization,
|
||||
})
|
||||
|
||||
@@ -290,14 +341,18 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
const estimated = MODEL_VRAM_ESTIMATES[estimateKey] ?? 165 * 1024 * 1024
|
||||
allocationToken = coordinator.requestAllocation(`kokoro-${quantization}`, estimated)
|
||||
|
||||
// Record manifest so consumers can inspect how the adapter resolved
|
||||
// device selection after fallback / WASM promotion.
|
||||
lastManifest = { quantization, device: (response.device ?? effectiveDevice) as string }
|
||||
|
||||
state = 'ready'
|
||||
updateInferenceStatus(modelStatusId, { state: 'ready', device: (response.device ?? device) as any })
|
||||
updateInferenceStatus(modelStatusId, { state: 'ready', device: (response.device ?? effectiveDevice) as any })
|
||||
onSuccess()
|
||||
if (!voices)
|
||||
throw new Error('Kokoro worker did not return voice metadata')
|
||||
return voices
|
||||
})
|
||||
}), { quantization, device }).catch((error) => {
|
||||
}), { quantization, device: effectiveDevice }).catch((error) => {
|
||||
handleWorkerError(error instanceof Error ? error : new Error(String(error)))
|
||||
throw error
|
||||
})
|
||||
@@ -364,6 +419,8 @@ export function createKokoroAdapter(): KokoroAdapter {
|
||||
getVoices,
|
||||
terminate: terminateAdapter,
|
||||
get state() { return state },
|
||||
get manifest() { return lastManifest },
|
||||
get deviceLossCount() { return deviceLossCount },
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -13,10 +13,10 @@ import { defaultPerfTracer } from '@proj-airi/stage-shared'
|
||||
import { Mutex } from 'async-mutex'
|
||||
|
||||
import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status'
|
||||
import { MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants'
|
||||
import { DEVICE_LOSS_WASM_THRESHOLD, MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants'
|
||||
import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator'
|
||||
import { LOAD_PRIORITY } from '../load-queue'
|
||||
import { createRequestId } from '../protocol'
|
||||
import { classifyDeviceLossReason, classifyError, createRequestId } from '../protocol'
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Types
|
||||
@@ -64,6 +64,15 @@ export interface WhisperAdapter {
|
||||
* Returns an unsubscribe function.
|
||||
*/
|
||||
onMessage: (handler: (event: WhisperEvent) => void) => () => void
|
||||
|
||||
/**
|
||||
* Snapshot of the last successful load, or null if never loaded.
|
||||
* `device` reflects what the worker actually used (post-fallback).
|
||||
*/
|
||||
readonly manifest: { device: string } | null
|
||||
|
||||
/** Number of WebGPU device-loss events observed by this adapter */
|
||||
readonly deviceLossCount: number
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -86,11 +95,26 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
|
||||
let errorListener: ((event: ErrorEvent) => void) | null = null
|
||||
const messageHandlers = new Set<(event: WhisperEvent) => void>()
|
||||
|
||||
// NOTICE: Device-loss resilience state. See kokoro.ts for rationale.
|
||||
let lastManifest: { device: string } | null = null
|
||||
let deviceLossCount = 0
|
||||
|
||||
const operationMutex = new Mutex()
|
||||
|
||||
function handleWorkerError(_event: ErrorEvent | Error): void {
|
||||
function handleWorkerError(event: ErrorEvent | Error): void {
|
||||
state = 'error'
|
||||
operationMutex.cancel()
|
||||
|
||||
const code = classifyError(event instanceof Error ? event : (event as ErrorEvent).error ?? event)
|
||||
if (code === 'DEVICE_LOST') {
|
||||
deviceLossCount++
|
||||
getGPUCoordinator().recordDeviceLoss({
|
||||
modelId: MODEL_NAMES.WHISPER,
|
||||
reason: classifyDeviceLossReason(event instanceof Error ? event : (event as ErrorEvent).error ?? event),
|
||||
occurredAt: Date.now(),
|
||||
})
|
||||
}
|
||||
|
||||
destroyWorker()
|
||||
scheduleRestart()
|
||||
}
|
||||
@@ -208,9 +232,20 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
|
||||
async function load(
|
||||
onProgress?: (p: ProgressPayload) => void,
|
||||
): Promise<void> {
|
||||
// NOTICE: Proactive WASM promotion after repeated device-loss events.
|
||||
// See kokoro.ts for rationale. Whisper always requests 'webgpu' from the
|
||||
// caller today, so we only check the promotion threshold.
|
||||
const requestedDevice = deviceLossCount >= DEVICE_LOSS_WASM_THRESHOLD ? 'wasm' : 'webgpu'
|
||||
if (requestedDevice === 'wasm') {
|
||||
console.warn(
|
||||
`[WhisperAdapter] ${deviceLossCount} device-loss events recorded, `
|
||||
+ `promoting load from webgpu to wasm.`,
|
||||
)
|
||||
}
|
||||
|
||||
return operationMutex.runExclusive(async () => {
|
||||
state = 'loading'
|
||||
updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'downloading', device: 'webgpu' })
|
||||
updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'downloading', device: requestedDevice as any })
|
||||
|
||||
return getLoadQueue().enqueue(MODEL_NAMES.WHISPER, LOAD_PRIORITY.ASR, async () => {
|
||||
const w = ensureWorker()
|
||||
@@ -230,7 +265,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
|
||||
}
|
||||
})
|
||||
|
||||
w.postMessage({ type: 'load-model', requestId, modelId: MODEL_NAMES.WHISPER, device: 'webgpu' })
|
||||
w.postMessage({ type: 'load-model', requestId, modelId: MODEL_NAMES.WHISPER, device: requestedDevice })
|
||||
|
||||
let readyResponse: any
|
||||
try {
|
||||
@@ -243,7 +278,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
|
||||
}
|
||||
|
||||
// Capture actual device reported by the worker (may fall back to WASM)
|
||||
const actualDevice = readyResponse?.device ?? 'webgpu'
|
||||
const actualDevice = readyResponse?.device ?? requestedDevice
|
||||
|
||||
// Track GPU memory allocation
|
||||
const coordinator = getGPUCoordinator()
|
||||
@@ -254,6 +289,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
|
||||
MODEL_VRAM_ESTIMATES[MODEL_NAMES.WHISPER] ?? 800 * 1024 * 1024,
|
||||
)
|
||||
|
||||
lastManifest = { device: actualDevice }
|
||||
state = 'ready'
|
||||
updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'ready', device: actualDevice })
|
||||
onSuccess()
|
||||
@@ -317,5 +353,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
|
||||
terminate: terminateAdapter,
|
||||
onMessage,
|
||||
get state() { return state },
|
||||
get manifest() { return lastManifest },
|
||||
get deviceLossCount() { return deviceLossCount },
|
||||
}
|
||||
}
|
||||
|
||||
@@ -53,3 +53,15 @@ export const MAX_RESTARTS = 3
|
||||
|
||||
/** Base delay in ms between restart attempts (multiplied by attempt number) */
|
||||
export const RESTART_DELAY_MS = 1_000
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Device loss resilience
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Number of WebGPU device-loss events an adapter tolerates before proactively
|
||||
* promoting subsequent loads to WASM. A single device loss may be transient
|
||||
* (driver reset, GPU process crash), but repeated losses indicate the WebGPU
|
||||
* path is unreliable on this device and WASM is safer.
|
||||
*/
|
||||
export const DEVICE_LOSS_WASM_THRESHOLD = 2
|
||||
|
||||
@@ -91,4 +91,51 @@ describe('gpuResourceCoordinator', () => {
|
||||
coordinator.requestAllocation('model', 700 * 1024 * 1024)
|
||||
expect(handler).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
describe('device loss telemetry', () => {
|
||||
it('should start with zero device-loss metrics', () => {
|
||||
const coordinator = createGPUResourceCoordinator(VRAM)
|
||||
const metrics = coordinator.getDeviceLossMetrics()
|
||||
|
||||
expect(metrics.totalCount).toBe(0)
|
||||
expect(metrics.byModel).toEqual({})
|
||||
expect(metrics.lastEvent).toBeNull()
|
||||
})
|
||||
|
||||
it('should aggregate device-loss events across models', () => {
|
||||
const coordinator = createGPUResourceCoordinator(VRAM)
|
||||
|
||||
coordinator.recordDeviceLoss({ modelId: 'kokoro', reason: 'unknown', occurredAt: 100 })
|
||||
coordinator.recordDeviceLoss({ modelId: 'kokoro', reason: 'unknown', occurredAt: 200 })
|
||||
coordinator.recordDeviceLoss({ modelId: 'whisper', reason: 'destroyed', occurredAt: 300 })
|
||||
|
||||
const metrics = coordinator.getDeviceLossMetrics()
|
||||
expect(metrics.totalCount).toBe(3)
|
||||
expect(metrics.byModel).toEqual({ kokoro: 2, whisper: 1 })
|
||||
expect(metrics.lastEvent).toEqual({ modelId: 'whisper', reason: 'destroyed', occurredAt: 300 })
|
||||
})
|
||||
|
||||
it('should notify subscribers on device-loss events', () => {
|
||||
const coordinator = createGPUResourceCoordinator(VRAM)
|
||||
const handler = vi.fn()
|
||||
coordinator.onDeviceLoss(handler)
|
||||
|
||||
const event = { modelId: 'kokoro', reason: 'unknown' as const, occurredAt: 100 }
|
||||
coordinator.recordDeviceLoss(event)
|
||||
|
||||
expect(handler).toHaveBeenCalledTimes(1)
|
||||
expect(handler).toHaveBeenCalledWith(event)
|
||||
})
|
||||
|
||||
it('should allow unsubscribing from device-loss events', () => {
|
||||
const coordinator = createGPUResourceCoordinator(VRAM)
|
||||
const handler = vi.fn()
|
||||
const unsub = coordinator.onDeviceLoss(handler)
|
||||
|
||||
unsub()
|
||||
coordinator.recordDeviceLoss({ modelId: 'x', reason: 'unknown', occurredAt: 1 })
|
||||
|
||||
expect(handler).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -7,8 +7,13 @@
|
||||
*
|
||||
* Emits memory pressure events when allocation nears the budget
|
||||
* so consumers can decide to unload LRU models or fall back to WASM.
|
||||
*
|
||||
* Also records device-loss telemetry so adapters can coordinate
|
||||
* cross-model WASM fallback decisions.
|
||||
*/
|
||||
|
||||
import type { DeviceLossReason } from './protocol'
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Types
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -31,6 +36,21 @@ export interface GPUResourceUsage {
|
||||
models: string[]
|
||||
}
|
||||
|
||||
export interface DeviceLossEvent {
|
||||
modelId: string
|
||||
reason: DeviceLossReason
|
||||
occurredAt: number
|
||||
}
|
||||
|
||||
export interface DeviceLossMetrics {
|
||||
/** Total device-loss events recorded across all models */
|
||||
totalCount: number
|
||||
/** Per-model device-loss counts */
|
||||
byModel: Record<string, number>
|
||||
/** Most recent event, or null if none recorded */
|
||||
lastEvent: DeviceLossEvent | null
|
||||
}
|
||||
|
||||
export interface GPUResourceCoordinator {
|
||||
/**
|
||||
* Request an allocation for a model.
|
||||
@@ -58,6 +78,22 @@ export interface GPUResourceCoordinator {
|
||||
* Returns an unsubscribe function.
|
||||
*/
|
||||
onMemoryPressure: (handler: (level: MemoryPressureLevel) => void) => () => void
|
||||
|
||||
/**
|
||||
* Record a WebGPU device-loss event. Adapters call this from their error
|
||||
* handlers when they detect a DEVICE_LOST error so the coordinator can
|
||||
* maintain cross-model telemetry.
|
||||
*/
|
||||
recordDeviceLoss: (event: DeviceLossEvent) => void
|
||||
|
||||
/** Get current device-loss telemetry across all models */
|
||||
getDeviceLossMetrics: () => DeviceLossMetrics
|
||||
|
||||
/**
|
||||
* Subscribe to device-loss events. Fired after `recordDeviceLoss()`.
|
||||
* Returns an unsubscribe function.
|
||||
*/
|
||||
onDeviceLoss: (handler: (event: DeviceLossEvent) => void) => () => void
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -78,6 +114,10 @@ export function createGPUResourceCoordinator(
|
||||
const budget = estimatedVRAM > 0 ? estimatedVRAM * BUDGET_SAFETY_FACTOR : Number.POSITIVE_INFINITY
|
||||
const allocations = new Map<string, AllocationToken>()
|
||||
const pressureHandlers = new Set<(level: MemoryPressureLevel) => void>()
|
||||
const deviceLossHandlers = new Set<(event: DeviceLossEvent) => void>()
|
||||
const deviceLossByModel = new Map<string, number>()
|
||||
let deviceLossTotal = 0
|
||||
let lastDeviceLossEvent: DeviceLossEvent | null = null
|
||||
|
||||
function getAllocated(): number {
|
||||
let total = 0
|
||||
@@ -154,6 +194,27 @@ export function createGPUResourceCoordinator(
|
||||
return () => pressureHandlers.delete(handler)
|
||||
}
|
||||
|
||||
function recordDeviceLoss(event: DeviceLossEvent): void {
|
||||
deviceLossTotal++
|
||||
deviceLossByModel.set(event.modelId, (deviceLossByModel.get(event.modelId) ?? 0) + 1)
|
||||
lastDeviceLossEvent = event
|
||||
for (const handler of deviceLossHandlers)
|
||||
handler(event)
|
||||
}
|
||||
|
||||
function getDeviceLossMetrics(): DeviceLossMetrics {
|
||||
return {
|
||||
totalCount: deviceLossTotal,
|
||||
byModel: Object.fromEntries(deviceLossByModel),
|
||||
lastEvent: lastDeviceLossEvent,
|
||||
}
|
||||
}
|
||||
|
||||
function onDeviceLoss(handler: (event: DeviceLossEvent) => void): () => void {
|
||||
deviceLossHandlers.add(handler)
|
||||
return () => deviceLossHandlers.delete(handler)
|
||||
}
|
||||
|
||||
return {
|
||||
requestAllocation,
|
||||
release,
|
||||
@@ -161,5 +222,8 @@ export function createGPUResourceCoordinator(
|
||||
getUsage,
|
||||
getLRUModel,
|
||||
onMemoryPressure,
|
||||
recordDeviceLoss,
|
||||
getDeviceLossMetrics,
|
||||
onDeviceLoss,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
import { classifyError, isRecoverable } from './protocol'
|
||||
import { classifyDeviceLossReason, classifyError, isRecoverable } from './protocol'
|
||||
|
||||
describe('classifyError', () => {
|
||||
it('should classify OOM errors', () => {
|
||||
@@ -13,6 +13,16 @@ describe('classifyError', () => {
|
||||
expect(classifyError(new Error('WebGPU device lost unexpectedly'))).toBe('DEVICE_LOST')
|
||||
})
|
||||
|
||||
it('should classify extended DEVICE_LOST patterns', () => {
|
||||
expect(classifyError(new Error('GPU device lost'))).toBe('DEVICE_LOST')
|
||||
expect(classifyError(new Error('GPUDevice was invalidated'))).toBe('DEVICE_LOST')
|
||||
expect(classifyError(new Error('GPUDevice is invalid'))).toBe('DEVICE_LOST')
|
||||
expect(classifyError(new Error('Device destroyed by user agent'))).toBe('DEVICE_LOST')
|
||||
expect(classifyError(new Error('GPU process crashed'))).toBe('DEVICE_LOST')
|
||||
expect(classifyError(new Error('GPU process lost'))).toBe('DEVICE_LOST')
|
||||
expect(classifyError(new Error('WebGPU device is invalid'))).toBe('DEVICE_LOST')
|
||||
})
|
||||
|
||||
it('should classify TIMEOUT errors', () => {
|
||||
expect(classifyError(new Error('operation timeout after 120s'))).toBe('TIMEOUT')
|
||||
})
|
||||
@@ -70,3 +80,27 @@ describe('isRecoverable', () => {
|
||||
expect(isRecoverable('UNKNOWN')).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe('classifyDeviceLossReason', () => {
|
||||
it('should return destroyed for GPUDeviceLostInfo-shaped object', () => {
|
||||
expect(classifyDeviceLossReason({ reason: 'destroyed', message: 'user requested' })).toBe('destroyed')
|
||||
})
|
||||
|
||||
it('should return unknown for non-destroyed structured reason', () => {
|
||||
expect(classifyDeviceLossReason({ reason: 'unknown', message: 'driver reset' })).toBe('unknown')
|
||||
})
|
||||
|
||||
it('should return destroyed when error message contains "destroyed"', () => {
|
||||
expect(classifyDeviceLossReason(new Error('Device destroyed by user agent'))).toBe('destroyed')
|
||||
})
|
||||
|
||||
it('should return unknown for generic device-loss messages', () => {
|
||||
expect(classifyDeviceLossReason(new Error('GPU device lost'))).toBe('unknown')
|
||||
expect(classifyDeviceLossReason(new Error('WebGPU device lost unexpectedly'))).toBe('unknown')
|
||||
})
|
||||
|
||||
it('should return unknown for non-device-loss inputs', () => {
|
||||
expect(classifyDeviceLossReason(new Error('out of memory'))).toBe('unknown')
|
||||
expect(classifyDeviceLossReason('random string')).toBe('unknown')
|
||||
})
|
||||
})
|
||||
|
||||
@@ -148,6 +148,26 @@ export function createRequestId(): string {
|
||||
return `req_${Date.now().toString(36)}_${(counter++).toString(36)}`
|
||||
}
|
||||
|
||||
// NOTICE: Patterns observed in WebGPU device loss errors across Chromium,
|
||||
// Firefox, Safari, and ONNX Runtime Web / transformers.js. Because we do not
|
||||
// own the GPUDevice (it is created internally by transformers.js / ORT-web),
|
||||
// we cannot attach a `device.lost` promise handler directly — string matching
|
||||
// on thrown errors is the only available detection signal.
|
||||
// References:
|
||||
// - https://gpuweb.github.io/gpuweb/#gpudevicelostinfo
|
||||
// - https://github.com/huggingface/transformers.js/issues/715
|
||||
const DEVICE_LOSS_PATTERNS = [
|
||||
'device was lost',
|
||||
'device lost',
|
||||
'gpu device lost',
|
||||
'gpudevice was invalidated',
|
||||
'gpudevice is invalid',
|
||||
'device destroyed',
|
||||
'gpu process crashed',
|
||||
'gpu process lost',
|
||||
'webgpu device is invalid',
|
||||
] as const
|
||||
|
||||
/**
|
||||
* Classify an unknown error into an `InferenceErrorCode`.
|
||||
* Used by worker adapters to normalise caught exceptions.
|
||||
@@ -162,7 +182,7 @@ export function classifyError(error: unknown, phase?: 'load' | 'inference'): Inf
|
||||
|
||||
if (lower.includes('out of memory') || lower.includes('allocation failed'))
|
||||
return 'OOM'
|
||||
if (lower.includes('device was lost') || lower.includes('device lost'))
|
||||
if (DEVICE_LOSS_PATTERNS.some(p => lower.includes(p)))
|
||||
return 'DEVICE_LOST'
|
||||
if (lower.includes('timeout'))
|
||||
return 'TIMEOUT'
|
||||
@@ -175,6 +195,31 @@ export function classifyError(error: unknown, phase?: 'load' | 'inference'): Inf
|
||||
return 'UNKNOWN'
|
||||
}
|
||||
|
||||
/** Reason classification for a device-loss event, following `GPUDeviceLostInfo.reason`. */
|
||||
export type DeviceLossReason = 'destroyed' | 'unknown'
|
||||
|
||||
/**
|
||||
* Best-effort classification of a device-loss reason from an error message
|
||||
* or a `GPUDeviceLostInfo`-shaped object. 'destroyed' implies intentional
|
||||
* termination (no recovery); 'unknown' implies a transient event that may
|
||||
* be recoverable via adapter restart or WASM fallback.
|
||||
*/
|
||||
export function classifyDeviceLossReason(error: unknown): DeviceLossReason {
|
||||
// Prefer structured info when available (some browsers attach GPUDeviceLostInfo)
|
||||
if (error && typeof error === 'object' && 'reason' in error) {
|
||||
const reason = (error as { reason?: unknown }).reason
|
||||
if (reason === 'destroyed')
|
||||
return 'destroyed'
|
||||
return 'unknown'
|
||||
}
|
||||
|
||||
const msg = error instanceof Error ? error.message : String(error)
|
||||
const lower = msg.toLowerCase()
|
||||
if (lower.includes('destroyed'))
|
||||
return 'destroyed'
|
||||
return 'unknown'
|
||||
}
|
||||
|
||||
/**
|
||||
* Determine whether an error code represents a potentially recoverable
|
||||
* condition. TIMEOUT and DEVICE_LOST may succeed on retry (e.g. with
|
||||
|
||||
Reference in New Issue
Block a user