fix(speech): keep a slow GPU model load on the GPU while it still makes progress (#40472)

The GPU stall watch gave the whole model load one 25 s budget, so a cold disk
or busy machine restarted a healthy load on the CPU and could miss the 60 s
start limit. The speech process now reports each load step, and the watch
restarts on it; a step that hangs for 25 s still moves to the CPU.
This commit is contained in:
Kelvin Amoaba
2026-11-09 20:49:37 +00:00
committed by GitHub
parent 81e5e5a9ef
commit 4e42cc2a3f
4 changed files with 40 additions and 3 deletions
+2
View File
@@ -15,6 +15,8 @@ export type SttRecognizerEvent =
| { type: 'chunk'; streamId: string; text: string; samples: number }
/** Pause-cut streams only: speech was heard that no chunk has committed yet. */
| { type: 'speech-pending'; streamId: string; pending: boolean }
/** A load step finished, so a slow load is told apart from a hung GPU. */
| { type: 'loading' }
export type SttRecognizerPost = (event: SttRecognizerEvent) => void
@@ -91,7 +91,12 @@ export async function createTranscribeCppRecognizer(
if (!gguf) {
throw new Error(`No .gguf model file in: ${init.files.join(', ')}`)
}
// Device setup (Metal shader library, Vulkan drivers) and reading the weights each report done,
// so a cold disk or a busy machine isn't taken for a GPU that hung.
await transcribe.initialize()
post({ type: 'loading' })
const model = await transcribe.TranscribeModel.load(join(init.modelDir, gguf))
post({ type: 'loading' })
// The child runs one call at a time, so every stream can share one session.
const session = model.createSession({ nThreads: performanceCoreThreads() })
const { languages } = model.capabilities
@@ -1,5 +1,5 @@
import { EventEmitter } from 'node:events'
import { describe, expect, it } from 'vitest'
import { describe, expect, it, vi } from 'vitest'
import { spawnSttWorkerWithCpuFallback } from './stt-worker-cpu-fallback'
class FakeChild extends EventEmitter {
@@ -115,6 +115,31 @@ describe('GPU speech process falling back to the CPU', () => {
expect(marks()).toBe(0)
})
it('keeps a slow GPU load that still finishes its steps, and moves a silent one to the CPU', () => {
vi.useFakeTimers()
try {
const { children, events } = setup()
const [gpu] = children
vi.advanceTimersByTime(20_000)
gpu.emit('message', { type: 'loading' })
vi.advanceTimersByTime(20_000)
gpu.emit('message', { type: 'loading' })
vi.advanceTimersByTime(20_000)
gpu.emit('message', { type: 'ready' })
expect(children).toHaveLength(1)
expect(events).toEqual([{ type: 'ready' }])
const silent = setup()
vi.advanceTimersByTime(20_000)
silent.children[0].emit('message', { type: 'loading' })
vi.advanceTimersByTime(25_000)
expect(silent.children).toHaveLength(2)
expect(silent.children[1].cpuOnly).toBe(true)
} finally {
vi.useRealTimers()
}
})
it('fails the worker when the CPU replacement cannot load the model mid-dictation', () => {
const { worker, children } = setup()
const exits: number[] = []
+7 -2
View File
@@ -17,7 +17,7 @@ type ChildMessage =
| { type: 'chunk'; streamId: string; text: string; samples: number }
| { type: 'speech-pending'; streamId: string; pending: boolean }
| {
type: 'ready' | 'stopped' | 'error' | 'partial' | 'final'
type: 'ready' | 'loading' | 'stopped' | 'error' | 'partial' | 'final'
streamId?: string
text?: string
error?: string
@@ -33,7 +33,8 @@ type StreamBacklog = {
pending: 'stop' | 'discard' | null
}
// Well under the session's 60 s start and stop timeouts, so a CPU replay still fits; covers first-run shader compiles.
// Per load step or decode. Well under the session's 60 s start and stop timeouts, so a CPU replay still fits;
// covers first-run shader compiles.
const GPU_STALL_TIMEOUT_MS = 25_000
/**
@@ -132,6 +133,10 @@ class SttCpuFallbackWorker extends EventEmitter implements SttWorkerProcess {
this.emit('exit', 1)
return
}
if (message.type === 'loading') {
this.watchForStall(true)
return
}
if (message.type === 'chunk') {
const { streamId, text, samples } = message
const stream = this.streams.get(streamId)