mirror of
https://github.com/stablyai/orca.git
synced 2026-11-14 01:47:24 +00:00
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:
@@ -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[] = []
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user