mara/hyperhive
mara/hyperhive
Watch
0
0
Fork
You've already forked hyperhive
 0
Code Activity

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

Browse source
This commit is contained in:
damocles 2026-11-12 08:07:38 +01:00 • committed by mara
commit f5e7caa2eb
2 changed files with 194 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

72
swarm-stt/src/engine.rs
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");
}
}

212
swarm-stt/src/main.rs
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());
}
}