swarm-stt: the daemon (queue connection, request/reply, busy limit, hourly purge) and its settings docs

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

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<u64> {
+        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<Whisper>, 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<String> {
+    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<ClientSecret> = 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"
+        ),
+    }
+}