diff --git a/docs/4-implementation/swarm/stt.md b/docs/4-implementation/swarm/stt.md index ce6627c25..05445e567 100644 --- a/docs/4-implementation/swarm/stt.md +++ b/docs/4-implementation/swarm/stt.md @@ -52,8 +52,14 @@ nowhere but the inbox of its own client, so replies leave through NATS reply rights. The service holds at most `MAX_WAITING` requests at once, whether running or -waiting for the engine's one slot, because each holds its audio in memory. It -answers a request past that with a `failed` response saying it's busy. It +waiting for the engine's one slot, because each holds its audio in memory, and +at most one per asking agent (the last token of the subject), so one agent +can't fill every slot. It answers a request past either limit with a `failed` +response saying so. `MAX_WAITING` times `RUN_TIMEOUT_SECS` must stay under the +15-minute reply deadline, because a request at the back of the line waits for +every run ahead of it; the daemon refuses to start with settings that break +that. A `failed` message tells the sender briefly what went wrong; the tools' +own output and the host's error text go to the log only. It removes expired transcripts every hour. ## Settings @@ -68,8 +74,8 @@ set them. | `CACHE_TTL_SECS` | one day | how long a transcript stays | | `MAX_BYTES` | 6 MiB | the largest audio file taken; the queue carries 8 MiB | | `MAX_AUDIO_SECS` | 300 | seconds of audio decoded; ffmpeg drops the rest | -| `RUN_TIMEOUT_SECS` | 300 | the longest one run may take | -| `MAX_WAITING` | 8 | requests held at once | +| `RUN_TIMEOUT_SECS` | 240 | the longest one run may take | +| `MAX_WAITING` | 3 | requests held at once | | `NATS_URL`, `OIDC_TOKEN_ENDPOINT`, `OIDC_CLIENT_ID`, `OIDC_CA_FILE` | required, the CA file optional | the queue connection | Two things the daemon doesn't do yet: sweep scratch directories a crash left diff --git a/swarm-stt/src/engine.rs b/swarm-stt/src/engine.rs index a8df1fa7f..54036fe68 100644 --- a/swarm-stt/src/engine.rs +++ b/swarm-stt/src/engine.rs @@ -14,6 +14,11 @@ pub trait Engine: Send + Sync { async fn transcribe(&self, audio: &[u8]) -> Result; } +/// What a sender is told when the service could not run the tools at all. +/// The detail (paths, OS errors, tool output) goes to the log, never to the +/// sender. +const COULD_NOT_START: &str = "the service could not start the transcription"; + /// `ffmpeg` decodes any container to 16 kHz mono WAV, then `whisper-cli` reads /// it with the configured model. #[derive(Debug, Clone)] @@ -42,14 +47,17 @@ impl Engine for Whisper { .to_owned(), ); } - let dir = tempfile_in(&self.scratch) + let dir = tempfile_in(&self.scratch).await.map_err(|e| { + tracing::warn!(error = %e, "making a scratch directory failed"); + COULD_NOT_START.to_owned() + })?; + // Dropping `run` on a timeout kills the child that was running. + let result = tokio::time::timeout(self.timeout, self.run(&dir, audio)) .await - .map_err(|e| format!("making a scratch directory: {e}"))?; - let result = match tokio::time::timeout(self.timeout, self.run(&dir, audio)).await { - Ok(result) => result, - // Dropping `run` kills the child that was running. - Err(_) => Err(format!("timed out after {}s", self.timeout.as_secs())), - }; + .unwrap_or_else(|_| { + tracing::warn!(secs = self.timeout.as_secs(), "a transcription timed out"); + Err("the transcription timed out".to_owned()) + }); let _ = tokio::fs::remove_dir_all(&dir).await; result } @@ -80,20 +88,24 @@ impl Whisper { .stderr(Stdio::piped()) .kill_on_drop(true) .spawn() - .map_err(|e| format!("starting ffmpeg: {e}"))?; + .map_err(|e| { + tracing::warn!(error = %e, "starting ffmpeg failed"); + COULD_NOT_START.to_owned() + })?; if let Some(mut stdin) = decode.stdin.take() { // ffmpeg may stop reading early (a cut at `-t`): a broken pipe is fine. let _ = stdin.write_all(audio).await; } - let decoded = decode - .wait_with_output() - .await - .map_err(|e| format!("waiting for ffmpeg: {e}"))?; + let decoded = decode.wait_with_output().await.map_err(|e| { + tracing::warn!(error = %e, "waiting for ffmpeg failed"); + COULD_NOT_START.to_owned() + })?; if !decoded.status.success() { - return Err(format!( - "ffmpeg could not decode the audio: {}", - String::from_utf8_lossy(&decoded.stderr).trim() - )); + tracing::warn!( + stderr = %String::from_utf8_lossy(&decoded.stderr).trim(), + "ffmpeg could not decode the audio" + ); + return Err("the audio could not be decoded".to_owned()); } let heard = Command::new(&self.whisper) .arg("-m") @@ -105,12 +117,16 @@ impl Whisper { .kill_on_drop(true) .output() .await - .map_err(|e| format!("starting whisper-cli: {e}"))?; + .map_err(|e| { + tracing::warn!(error = %e, "starting whisper-cli failed"); + COULD_NOT_START.to_owned() + })?; if !heard.status.success() { - return Err(format!( - "whisper-cli failed: {}", - String::from_utf8_lossy(&heard.stderr).trim() - )); + tracing::warn!( + stderr = %String::from_utf8_lossy(&heard.stderr).trim(), + "whisper-cli failed" + ); + return Err("the transcription failed".to_owned()); } Ok(String::from_utf8_lossy(&heard.stdout).trim().to_owned()) } @@ -221,24 +237,18 @@ mod tests { } #[tokio::test] - async fn an_ffmpeg_failure_is_reported_with_its_stderr() { + async fn an_ffmpeg_failure_tells_the_sender_briefly_and_keeps_the_detail_out() { let dir = tempfile::tempdir().unwrap(); let engine = whisper(dir.path(), "echo 'bad data' >&2; exit 1", WHISPER_ECHO); let err = engine.transcribe(b"OggSx").await.unwrap_err(); - assert!( - err.contains("could not decode") && err.contains("bad data"), - "{err}" - ); + assert_eq!(err, "the audio could not be decoded"); } #[tokio::test] - async fn a_whisper_failure_is_reported_with_its_stderr() { + async fn a_whisper_failure_tells_the_sender_briefly_and_keeps_the_detail_out() { let dir = tempfile::tempdir().unwrap(); let engine = whisper(dir.path(), FFMPEG_OK, "echo 'no model' >&2; exit 2"); let err = engine.transcribe(b"OggSx").await.unwrap_err(); - assert!( - err.contains("whisper-cli failed") && err.contains("no model"), - "{err}" - ); + assert_eq!(err, "the transcription failed"); } } diff --git a/swarm-stt/src/main.rs b/swarm-stt/src/main.rs index e78ca81bc..50d6542b3 100644 --- a/swarm-stt/src/main.rs +++ b/swarm-stt/src/main.rs @@ -6,8 +6,9 @@ //! runs and the engine. This file only reads the settings and carries //! requests between the queue and the service. +use std::collections::HashSet; use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -33,6 +34,82 @@ const QUEUE_GROUP: &str = "swarm-stt"; /// How often expired transcripts are removed. const PURGE_EVERY: Duration = Duration::from_hours(1); +type Var<'a> = &'a dyn Fn(&str) -> Result; +type Number<'a> = &'a dyn Fn(&str, u64) -> Result; + +/// The service and the limits around it, from the environment. +struct Settings { + service: Arc>, + cache: Cache, + waiting: Arc, +} + +fn settings(var: Var<'_>, number: Number<'_>) -> Result { + let state = PathBuf::from(var("STATE_DIR")?); + let cache = Cache::new( + state.join("cache"), + Duration::from_secs(number("CACHE_TTL_SECS", 24 * 60 * 60)?), + ); + let engine_timeout_secs = number("RUN_TIMEOUT_SECS", 240)?; + let engine = Whisper { + ffmpeg: PathBuf::from(var("FFMPEG")?), + whisper: PathBuf::from(var("WHISPER_CLI")?), + model: PathBuf::from(var("MODEL")?), + scratch: state.join("scratch"), + max_secs: number("MAX_AUDIO_SECS", 300)?, + timeout: Duration::from_secs(engine_timeout_secs), + }; + let max_bytes = usize::try_from(number("MAX_BYTES", 6 * 1024 * 1024)?) + .context("SWARM_STT_MAX_BYTES does not fit in memory")?; + let service = Arc::new(Service::new(engine, cache.clone(), Config { max_bytes })); + // Requests held at once, running or waiting for the engine's one slot: + // each holds its audio in memory, so past this the service says it is busy. + let max_waiting = number("MAX_WAITING", 3)?; + let worst_wait = Duration::from_secs(max_waiting.saturating_mul(engine_timeout_secs)); + if worst_wait >= swarm_wire::stt::REPLY_DEADLINE { + bail!( + "SWARM_STT_MAX_WAITING ({max_waiting}) times SWARM_STT_RUN_TIMEOUT_SECS \ + ({engine_timeout_secs}) must stay under the {}s reply deadline, or a queued \ + request is answered after its sender has given up", + swarm_wire::stt::REPLY_DEADLINE.as_secs() + ); + } + let waiting = Arc::new(Semaphore::new( + usize::try_from(max_waiting).context("SWARM_STT_MAX_WAITING is too large")?, + )); + Ok(Settings { + service, + cache, + waiting, + }) +} + +/// The queue connection, as the service's own OIDC client. +async fn connect_queue(var: Var<'_>) -> Result { + let client_id = var("OIDC_CLIENT_ID")?; + let secret = + swarm_secret_client::read_with_startup_retry("reading the queue client secret", || { + client_secret(&client_id) + }) + .await?; + let reread_id = client_id.clone(); + let secret = StoredSecret::new(client_id.clone(), secret, move || { + let client_id = reread_id.clone(); + async move { client_secret(&client_id).await.map_err(Into::into) } + }); + swarm_queue_client::connect(QueueConfig { + url: var("NATS_URL")?, + token_endpoint: var("OIDC_TOKEN_ENDPOINT")?, + client_id, + client_secret: ClientSecret::Stored(secret), + ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE")) + .ok() + .map(PathBuf::from), + }) + .await + .map_err(Into::into) +} + #[tokio::main] async fn main() -> Result<()> { hive_log::init(); @@ -49,49 +126,13 @@ async fn main() -> Result<()> { } }; - let state = PathBuf::from(var("STATE_DIR")?); - let cache = Cache::new( - state.join("cache"), - Duration::from_secs(number("CACHE_TTL_SECS", 24 * 60 * 60)?), - ); - let engine = Whisper { - ffmpeg: PathBuf::from(var("FFMPEG")?), - whisper: PathBuf::from(var("WHISPER_CLI")?), - model: PathBuf::from(var("MODEL")?), - scratch: state.join("scratch"), - max_secs: number("MAX_AUDIO_SECS", 300)?, - timeout: Duration::from_secs(number("RUN_TIMEOUT_SECS", 300)?), - }; - let max_bytes = usize::try_from(number("MAX_BYTES", 6 * 1024 * 1024)?) - .context("SWARM_STT_MAX_BYTES does not fit in memory")?; - let service = Arc::new(Service::new(engine, cache.clone(), Config { max_bytes })); - // Requests held at once, running or waiting for the engine's one slot: - // each holds its audio in memory, so past this the service says it is busy. - let waiting = Arc::new(Semaphore::new( - usize::try_from(number("MAX_WAITING", 8)?).context("SWARM_STT_MAX_WAITING is too large")?, - )); - - let client_id = var("OIDC_CLIENT_ID")?; - let secret = - swarm_secret_client::read_with_startup_retry("reading the queue client secret", || { - client_secret(&client_id) - }) - .await?; - let reread_id = client_id.clone(); - let secret = StoredSecret::new(client_id.clone(), secret, move || { - let client_id = reread_id.clone(); - async move { client_secret(&client_id).await.map_err(Into::into) } - }); - let queue = swarm_queue_client::connect(QueueConfig { - url: var("NATS_URL")?, - token_endpoint: var("OIDC_TOKEN_ENDPOINT")?, - client_id, - client_secret: ClientSecret::Stored(secret), - ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE")) - .ok() - .map(PathBuf::from), - }) - .await?; + let Settings { + service, + cache, + waiting, + } = settings(&var, &number)?; + let per_agent: Arc>> = Arc::default(); + let queue = connect_queue(&var).await?; tokio::spawn(async move { loop { @@ -113,20 +154,38 @@ async fn main() -> Result<()> { ); while let Some(request) = requests.next().await { let (queue, service) = (queue.clone(), Arc::clone(&service)); - let permit = Arc::clone(&waiting).try_acquire_owned().ok(); + // Held at once overall, and one per asking agent, so one agent + // cannot fill every slot. The agent is the subject's last token. + let agent = request + .subject + .as_str() + .rsplit('.') + .next() + .unwrap_or_default() + .to_owned(); + let turned_away = match ( + Arc::clone(&waiting).try_acquire_owned(), + take_slot(&per_agent, &agent), + ) { + (Ok(permit), Some(slot)) => Ok((permit, slot)), + (Ok(_), None) => Err("you already have a transcription running; wait for it"), + (Err(_), _) => Err("the transcription service is busy, try again shortly"), + }; tokio::spawn(async move { let Some(reply) = request.reply.clone() else { tracing::warn!(subject = %request.subject, "a request with no reply subject"); return; }; - let response = if permit.is_some() { - answer(&service, &request.payload).await - } else { - Response::Failed { - message: "the transcription service is busy, try again shortly".to_owned(), + let response = match turned_away { + Ok(held) => { + let response = answer(&service, &request.payload).await; + drop(held); + response } + Err(message) => Response::Failed { + message: message.to_owned(), + }, }; - drop(permit); match serde_json::to_vec(&response) { Ok(body) => { if let Err(e) = queue.publish(reply, body.into()).await { @@ -140,6 +199,33 @@ async fn main() -> Result<()> { bail!("the request subscription ended") } +/// One agent's place among the requests in flight; dropping it frees the place. +struct Slot { + agents: Arc>>, + agent: String, +} + +impl Drop for Slot { + fn drop(&mut self) { + self.agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&self.agent); + } +} + +/// `agent`'s place, or `None` when it already has a request in flight. +fn take_slot(agents: &Arc>>, agent: &str) -> Option { + let fresh = agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(agent.to_owned()); + fresh.then(|| Slot { + agents: Arc::clone(agents), + agent: agent.to_owned(), + }) +} + /// The wire answer to one request's audio. async fn answer(service: &Service, audio: &[u8]) -> Response { match service.transcribe(audio).await { @@ -176,3 +262,27 @@ async fn client_secret(client_id: &str) -> Result { ), } } + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::sync::{Arc, Mutex}; + + use super::take_slot; + + #[test] + fn an_agent_holds_one_slot_at_a_time_and_dropping_it_frees_it() { + let agents: Arc>> = Arc::default(); + let first = take_slot(&agents, "iris").expect("a free agent gets a slot"); + assert!( + take_slot(&agents, "iris").is_none(), + "a second request is turned away" + ); + assert!( + take_slot(&agents, "argus").is_some(), + "another agent is not" + ); + drop(first); + assert!(take_slot(&agents, "iris").is_some()); + } +} diff --git a/swarm-wire/src/stt.rs b/swarm-wire/src/stt.rs index 0e0d5a525..5b4a97a46 100644 --- a/swarm-wire/src/stt.rs +++ b/swarm-wire/src/stt.rs @@ -29,9 +29,12 @@ pub fn subject_wildcard() -> String { } /// How long a sender waits for the answer, and the queue lets the service -/// reply: the longest run, plus the time a request may wait for the engine's -/// one slot. -pub const REPLY_DEADLINE: Duration = Duration::from_mins(10); +/// reply. The service holds a request at most `max_waiting` runs deep, each +/// at most its run timeout, so a request that waited for every slot ahead of +/// it still answers inside this; the daemon refuses settings where +/// `max_waiting * run_timeout` reaches it. Past the deadline the sender has +/// given up and the queue drops the reply. +pub const REPLY_DEADLINE: Duration = Duration::from_mins(15); /// The service's answer to one request. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]