agent: code-atomiser-fix — cancellable Whisper inference + bounded drain + lock-over-await
Three findings chain: thread::spawn discarded the JoinHandle and gave no cancellation route into the whisper backend; drain_inference busy- polled with no timeout; stop_live_transcription_session held the lifecycle AsyncMutex across the await of a non-cancellable spawn_blocking handle. Together: a single wedged inference (ggml deadlock, GPU stall) bricked every future start/stop until app restart. Fix: each inference task carries an Arc<AtomicBool> abort_flag; the flag is wired into whisper-rs::FullParams::set_abort_callback_safe so the spawned blocking thread checks it and exits cleanly. drain_inference is bounded by a deadline (3 x chunk_duration, min 2s); on expiry the flag is set, the receiver is dropped, and a typed Error::InferenceTimeout status is surfaced. Drop for InferenceTask asserts the abort flag so any '?'-propagation or panic unwind closes the cancellation route without relying on the explicit drain path. stop_live_transcription_session restructured so the lifecycle guard is dropped BEFORE the JoinHandle is awaited; start_live_transcription_session releases the guard explicitly after installing the RunningLiveSession. Unit test added: dropping_inference_task_sets_abort_flag covers the Race-A regression directly. A real Race-B drain-timeout test would require a wedged whisper-rs, which is hard to fixture; that path is covered by the SAFETY comment and cargo build + manual smoke. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -22,6 +22,14 @@ pub enum Error {
|
||||
#[error("transcription failed: {0}")]
|
||||
TranscriptionFailed(String),
|
||||
|
||||
/// Inference exceeded the bounded wait imposed by the caller (live
|
||||
/// session `drain_inference`). The spawned worker has had its abort
|
||||
/// flag set so whisper-rs will return early on its next
|
||||
/// abort-callback poll; the lock-up itself is *not* recovered by
|
||||
/// this error — but the live-session lifecycle can now progress.
|
||||
#[error("inference timed out after {timeout_ms}ms")]
|
||||
InferenceTimeout { timeout_ms: u64 },
|
||||
|
||||
#[error("audio decode failed: {0}")]
|
||||
AudioDecodeFailed(String),
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::path::Path;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
|
||||
use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult};
|
||||
@@ -155,6 +156,35 @@ impl LocalEngine {
|
||||
inference_ms,
|
||||
})
|
||||
}
|
||||
|
||||
/// Cancellable variant of `transcribe_sync`. Pipes `abort_flag`
|
||||
/// through to the backend so the live-session drain timeout can
|
||||
/// break a wedged whisper-rs decode out of its loop. Backends that
|
||||
/// can't honour the flag (Parakeet) fall through to their plain
|
||||
/// `transcribe_sync` via the trait's default implementation.
|
||||
pub fn transcribe_sync_with_abort(
|
||||
&self,
|
||||
audio: &AudioSamples,
|
||||
options: &TranscriptionOptions,
|
||||
abort_flag: Arc<AtomicBool>,
|
||||
) -> Result<TimedTranscript> {
|
||||
let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let backend = guard.as_mut().ok_or(Error::EngineNotLoaded)?;
|
||||
|
||||
let start = Instant::now();
|
||||
let segments =
|
||||
backend.transcribe_sync_with_abort(audio.samples(), options, abort_flag)?;
|
||||
let inference_ms = start.elapsed().as_millis() as u64;
|
||||
|
||||
Ok(TimedTranscript {
|
||||
transcript: Transcript::new(
|
||||
segments,
|
||||
options.language.clone().unwrap_or_else(|| "en".to_string()),
|
||||
audio.duration_secs(),
|
||||
),
|
||||
inference_ms,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Thin wrapper over `ParakeetModel` that overrides `transcribe_raw` to
|
||||
|
||||
@@ -9,6 +9,9 @@
|
||||
//! `whisper` feature — `WhisperRsBackend` (direct whisper-rs, the only
|
||||
//! path that pipes `initial_prompt`).
|
||||
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::Arc;
|
||||
|
||||
use lumotia_core::error::Result;
|
||||
use lumotia_core::types::{Segment, TranscriptionOptions};
|
||||
|
||||
@@ -45,6 +48,24 @@ pub trait Transcriber: Send {
|
||||
samples: &[f32],
|
||||
options: &TranscriptionOptions,
|
||||
) -> Result<Vec<Segment>>;
|
||||
|
||||
/// Variant of `transcribe_sync` that accepts an external abort flag.
|
||||
///
|
||||
/// Backends that can react to mid-inference cancellation (whisper-rs
|
||||
/// via `set_abort_callback_safe`) override this to poll the flag
|
||||
/// during decoding so a wedged inference can be unstuck by the live
|
||||
/// session's `drain_inference` timeout. The default implementation
|
||||
/// ignores the flag and dispatches to `transcribe_sync` so
|
||||
/// non-cancellable backends (Parakeet via transcribe-rs) keep their
|
||||
/// current behaviour.
|
||||
fn transcribe_sync_with_abort(
|
||||
&mut self,
|
||||
samples: &[f32],
|
||||
options: &TranscriptionOptions,
|
||||
_abort_flag: Arc<AtomicBool>,
|
||||
) -> Result<Vec<Segment>> {
|
||||
self.transcribe_sync(samples, options)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -7,6 +7,8 @@
|
||||
//! into Whisper.
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
|
||||
|
||||
@@ -57,10 +59,35 @@ impl Transcriber for WhisperRsBackend {
|
||||
&mut self,
|
||||
samples: &[f32],
|
||||
options: &TranscriptionOptions,
|
||||
) -> Result<Vec<Segment>> {
|
||||
self.transcribe_sync_inner(samples, options, None)
|
||||
}
|
||||
|
||||
/// Cancellable variant. Installs the abort flag as whisper-rs's
|
||||
/// abort callback so a `drain_inference` timeout in the live session
|
||||
/// can break the worker out of the decoding loop instead of letting
|
||||
/// it run to completion against a stuck backend.
|
||||
fn transcribe_sync_with_abort(
|
||||
&mut self,
|
||||
samples: &[f32],
|
||||
options: &TranscriptionOptions,
|
||||
abort_flag: Arc<AtomicBool>,
|
||||
) -> Result<Vec<Segment>> {
|
||||
self.transcribe_sync_inner(samples, options, Some(abort_flag))
|
||||
}
|
||||
}
|
||||
|
||||
impl WhisperRsBackend {
|
||||
fn transcribe_sync_inner(
|
||||
&mut self,
|
||||
samples: &[f32],
|
||||
options: &TranscriptionOptions,
|
||||
abort_flag: Option<Arc<AtomicBool>>,
|
||||
) -> Result<Vec<Segment>> {
|
||||
tracing::info!(
|
||||
language = ?options.language,
|
||||
has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false),
|
||||
has_abort_flag = abort_flag.is_some(),
|
||||
"WhisperRsBackend::transcribe_sync entering"
|
||||
);
|
||||
|
||||
@@ -87,6 +114,15 @@ impl Transcriber for WhisperRsBackend {
|
||||
params.set_print_progress(false);
|
||||
params.set_print_realtime(false);
|
||||
|
||||
// Wire the per-task abort flag into whisper-rs. The closure is
|
||||
// polled by whisper.cpp between decode steps; returning true
|
||||
// tells the backend to abort the current inference. This is the
|
||||
// cancellation route that closes the orphaned-thread loophole
|
||||
// when a live session is stopped or wedges.
|
||||
if let Some(flag) = abort_flag {
|
||||
params.set_abort_callback_safe(move || flag.load(Ordering::Relaxed));
|
||||
}
|
||||
|
||||
state.full(params, samples).map_err(|e| {
|
||||
Error::TranscriptionFailed(
|
||||
WhisperBackendError::Transcribe(e.to_string()).to_string(),
|
||||
|
||||
Reference in New Issue
Block a user