Watch
0
0
Fork
You've already forked hyperhive
0

swarm-stt: reply deadline agrees with the limits, one request per agent, short errors to the sender

This commit is contained in:
damocles 2026-11-12 08:07:38 +01:00 • committed by mara
commit f5e7caa2eb
4 changed files with 194 additions and 65 deletions

View file

@ -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

View file

@ -14,6 +14,11 @@ pub trait Engine: Send + Sync {
async fn transcribe(&self, audio: &[u8]) -> Result<String, String>;
}
/// 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");
}
}

View file

@ -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<String>;
type Number<'a> = &'a dyn Fn(&str, u64) -> Result<u64>;
/// The service and the limits around it, from the environment.
struct Settings {
service: Arc<Service<Whisper>>,
cache: Cache,
waiting: Arc<Semaphore>,
}
fn settings(var: Var<'_>, number: Number<'_>) -> Result<Settings> {
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<async_nats::Client> {
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<Mutex<HashSet<String>>> = 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<Mutex<HashSet<String>>>,
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<Mutex<HashSet<String>>>, agent: &str) -> Option<Slot> {
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<Whisper>, audio: &[u8]) -> Response {
match service.transcribe(audio).await {
@ -176,3 +262,27 @@ async fn client_secret(client_id: &str) -> Result<String> {
),
}
}
#[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<Mutex<HashSet<String>>> = 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());
}
}

View file

@ -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)]