diff --git a/Cargo.lock b/Cargo.lock index 0325825db..d20316d18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5368,8 +5368,16 @@ dependencies = [ name = "swarm-stt" version = "0.1.0" dependencies = [ + "anyhow", + "async-nats", "async-trait", + "futures-util", + "hive-log", + "serde_json", "sha2 0.11.1", + "swarm-queue-client", + "swarm-secret-client", + "swarm-wire", "tempfile", "thiserror 2.0.21", "tokio", diff --git a/docs/4-implementation/swarm/stt.md b/docs/4-implementation/swarm/stt.md index 903c883d1..ce6627c25 100644 --- a/docs/4-implementation/swarm/stt.md +++ b/docs/4-implementation/swarm/stt.md @@ -39,10 +39,39 @@ way to open them. The service kills a run that outlasts its timeout (decode and together) and fails the request, so one hung child can't hold the single engine slot. A scratch directory per run holds the WAV and goes away with the run. -## Settings the daemon will expose +## The daemon -The model (default `small`), the cache TTL (default one day), the size, -length and run-time caps. The daemon change also owns a bound on how many -requests may wait for the engine slot, a sweep of scratch directories a crash -left behind, and whisper's thread count. The design and the choices behind them are in the discussion on -issue 6383. +`swarm-stt` connects to the queue as its own OIDC client, with its secret read +from the store under its own certificate, as the notification router does. It +subscribes to `$SWARM.stt.transcribe.*` in the queue group `swarm-stt` and +answers each request on its reply subject. The body of a request is the raw +audio. The answer is JSON: `{"status":"ok","text":…,"cached":…}` or +`{"status":"failed","message":…}` (`swarm_wire::stt`). An agent may publish +only on its own `$SWARM.stt.transcribe.`, and the service may publish +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 +removes expired transcripts every hour. + +## Settings + +All are environment variables with the prefix `SWARM_STT_`; the nix module will +set them. + +| variable | default | meaning | +| ------------------------------------------------------------------- | ------------------------------ | ----------------------------------------------------- | +| `STATE_DIR` | required | holds `cache/` and `scratch/` | +| `FFMPEG`, `WHISPER_CLI`, `MODEL` | required | the two binaries and the ggml model file | +| `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 | +| `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 +behind, and set whisper's thread count. The issue 6383 discussion has the +choices behind the defaults. diff --git a/swarm-stt/Cargo.toml b/swarm-stt/Cargo.toml index a65bc3abc..5dc411632 100644 --- a/swarm-stt/Cargo.toml +++ b/swarm-stt/Cargo.toml @@ -3,11 +3,27 @@ name = "swarm-stt" version.workspace = true edition.workspace = true +[[bin]] +name = "swarm-stt" +path = "src/main.rs" + [dependencies] +anyhow.workspace = true +async-nats.workspace = true # The `Engine` trait: a fake in the tests, whisper in the daemon. async-trait.workspace = true +# Draining the subscription. +futures-util.workspace = true +hive-log.workspace = true +serde_json.workspace = true # The SHA-256 of the audio, which keys the cache and the in-flight map. sha2.workspace = true +# The queue connect with its OIDC client credentials. +swarm-queue-client.workspace = true +# The client secret, read with the service's own store leaf. +swarm-secret-client.workspace = true +# The subject, request and response of the wire. +swarm-wire.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true diff --git a/swarm-stt/src/main.rs b/swarm-stt/src/main.rs new file mode 100644 index 000000000..e78ca81bc --- /dev/null +++ b/swarm-stt/src/main.rs @@ -0,0 +1,178 @@ +//! The swarm's speech-to-text service. +//! +//! It connects to the queue as its own OIDC client, hears every agent's +//! request subject (`swarm_wire::stt`), and answers each audio file with its +//! transcript. [`swarm_stt::Service`] does the work: the cache, the shared +//! runs and the engine. This file only reads the settings and carries +//! requests between the queue and the service. + +use std::path::PathBuf; +use std::sync::Arc; + +use std::time::Duration; + +use anyhow::{Context as _, Result, bail}; +use futures_util::StreamExt as _; +use swarm_queue_client::{ClientSecret, QueueConfig}; +use swarm_secret_client::stored_secret::StoredSecret; +use swarm_stt::engine::Whisper; +use swarm_stt::{Config, Error, Service, cache::Cache}; +use swarm_wire::stt::{Response, subject_wildcard}; +use tokio::sync::Semaphore; + +/// The cert-auth role this service logs in to the store as. `swarm-bao.nix` +/// writes it, matching the leaf's subject. +const CERT_ROLE: &str = "swarm-stt"; + +/// The environment prefix `swarm-stt.nix` sets. +const PREFIX: &str = "SWARM_STT"; + +/// The queue group every instance joins, so a request reaches one of them. +const QUEUE_GROUP: &str = "swarm-stt"; + +/// How often expired transcripts are removed. +const PURGE_EVERY: Duration = Duration::from_hours(1); + +#[tokio::main] +async fn main() -> Result<()> { + hive_log::init(); + swarm_secret_client::install_crypto_provider(); + let var = |name: &str| { + let full = format!("{PREFIX}_{name}"); + std::env::var(&full).with_context(|| format!("{full} is unset")) + }; + let number = |name: &str, default: u64| -> Result { + let full = format!("{PREFIX}_{name}"); + match std::env::var(&full) { + Ok(v) => v.parse().with_context(|| format!("{full} is not a number")), + Err(_) => Ok(default), + } + }; + + 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?; + + tokio::spawn(async move { + loop { + let removed = cache.purge().await; + if removed > 0 { + tracing::info!(removed, "removed expired transcripts"); + } + tokio::time::sleep(PURGE_EVERY).await; + } + }); + + let mut requests = queue + .queue_subscribe(subject_wildcard(), QUEUE_GROUP.to_owned()) + .await + .context("subscribing to the request subject")?; + tracing::info!( + subject = subject_wildcard(), + "swarm-stt: answering requests" + ); + while let Some(request) = requests.next().await { + let (queue, service) = (queue.clone(), Arc::clone(&service)); + let permit = Arc::clone(&waiting).try_acquire_owned().ok(); + 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(), + } + }; + drop(permit); + match serde_json::to_vec(&response) { + Ok(body) => { + if let Err(e) = queue.publish(reply, body.into()).await { + tracing::warn!(error = %e, "publishing a reply failed"); + } + } + Err(e) => tracing::error!(error = %e, "encoding a reply failed"), + } + }); + } + bail!("the request subscription ended") +} + +/// The wire answer to one request's audio. +async fn answer(service: &Service, audio: &[u8]) -> Response { + match service.transcribe(audio).await { + Ok(reply) => Response::Ok { + text: reply.text, + cached: reply.cached, + }, + Err(Error::Engine(message)) => { + tracing::warn!(%message, "a transcription failed"); + Response::Failed { message } + } + Err(e) => Response::Failed { + message: e.to_string(), + }, + } +} + +async fn client_secret(client_id: &str) -> Result { + use swarm_secret_client::builder::{ClientSecret, client_secret_path}; + + let path = client_secret_path(client_id)?; + let store = swarm_secret_client::SecretStore::from_env(CERT_ROLE) + .await + .context("logging in to the secret store")?; + let stored: Option = store + .read_optional(&path) + .await + .with_context(|| format!("reading {path}"))?; + match stored { + Some(s) if !s.value.trim().is_empty() => Ok(s.value.trim().to_owned()), + _ => bail!( + "nothing at {path}: authelia's mint, in its container, has not \ + stored this service's client secret" + ), + } +}