#![forbid(unsafe_code)]
pub mod admission;
pub mod audio;
pub mod health;
use std::{
env, fmt,
ops::Range,
sync::{
Arc, OnceLock,
atomic::{AtomicBool, Ordering},
mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TrySendError},
},
thread,
time::{Duration, Instant},
};
use asupersync::runtime::{Runtime, RuntimeBuilder};
pub const SCAFFOLD_REVISION: u8 = 2;
const DEFAULT_QUEUE_CAPACITY: usize = 8;
const DEFAULT_SYNTHESIS_BUDGET: Duration = Duration::from_secs(30);
const DEFAULT_ENROLL_BUDGET: Duration = Duration::from_secs(30);
const BACKPRESSURE_POLL: Duration = Duration::from_millis(1);
const DEFAULT_SYNTHESIS_FRAME_BUDGET: Duration = Duration::from_secs(8);
const DEBUG_BUILD_SLOWDOWN: u32 = 32;
pub fn process_engine_config() -> EngineConfig {
static CONFIG: OnceLock<EngineConfig> = OnceLock::new();
CONFIG.get_or_init(EngineConfig::from_environment).clone()
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EngineConfig {
pub stream_queue_capacity: usize,
pub synthesis_stage_budget: Duration,
pub synthesis_frame_budget: Duration,
pub enroll_stage_budget: Duration,
pub admission: admission::AdmissionPolicy,
}
impl Default for EngineConfig {
fn default() -> Self {
let slowdown = build_profile_slowdown();
Self {
stream_queue_capacity: DEFAULT_QUEUE_CAPACITY,
synthesis_stage_budget: DEFAULT_SYNTHESIS_BUDGET * slowdown,
synthesis_frame_budget: DEFAULT_SYNTHESIS_FRAME_BUDGET * slowdown,
enroll_stage_budget: DEFAULT_ENROLL_BUDGET,
admission: admission::AdmissionPolicy::default(),
}
}
}
const fn build_profile_slowdown() -> u32 {
if cfg!(debug_assertions) {
DEBUG_BUILD_SLOWDOWN
} else {
1
}
}
impl EngineConfig {
fn from_environment() -> Self {
let mut config = Self::default();
config.synthesis_stage_budget = stage_budget_from_environment(
"FTTS_STAGE_BUDGET_SYNTHESIS_MS",
config.synthesis_stage_budget,
);
config.synthesis_frame_budget = stage_budget_from_environment(
"FTTS_STAGE_BUDGET_FRAME_MS",
config.synthesis_frame_budget,
);
config.enroll_stage_budget = stage_budget_from_environment(
"FTTS_STAGE_BUDGET_ENROLL_MS",
config.enroll_stage_budget,
);
config.admission.budget_bytes = positive_u64_from_environment("FTTS_MEMORY_BUDGET_MB")
.and_then(|megabytes| megabytes.checked_mul(1024 * 1024))
.unwrap_or(config.admission.budget_bytes);
if let Some(max_frames) = positive_u64_from_environment("FTTS_MAX_FRAMES") {
config.admission.max_new_tokens = max_frames;
config.admission.heuristic_eos_backstop = false;
}
config
}
fn validate(&self) -> Result<(), EngineError> {
if self.stream_queue_capacity == 0 {
return Err(EngineError::InvalidConfiguration(
"stream queue capacity must be greater than zero",
));
}
if self.synthesis_stage_budget.is_zero()
|| self.synthesis_frame_budget.is_zero()
|| self.enroll_stage_budget.is_zero()
{
return Err(EngineError::InvalidConfiguration(
"stage budgets must be greater than zero",
));
}
Ok(())
}
}
fn positive_u64_from_environment(name: &str) -> Option<u64> {
env::var(name)
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|value| *value > 0)
}
fn stage_budget_from_environment(name: &str, fallback: Duration) -> Duration {
env::var(name)
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|milliseconds| *milliseconds > 0)
.map(Duration::from_millis)
.unwrap_or(fallback)
}
#[derive(Clone, Debug, Default)]
pub struct CancellationToken {
cancelled: Arc<AtomicBool>,
}
impl CancellationToken {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn cancel(&self) {
self.cancelled.store(true, Ordering::Release);
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.cancelled.load(Ordering::Acquire)
}
pub fn checkpoint(&self) -> Result<(), EngineError> {
if self.is_cancelled() {
Err(EngineError::Cancelled)
} else {
Ok(())
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StreamKind {
Pcm,
Events,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PcmPacket {
pub frame_count: u8,
pub samples: Vec<i16>,
}
pub struct StreamQueues {
pub pcm: BoundedSender<PcmPacket>,
pub pcm_receiver: BoundedReceiver<PcmPacket>,
pub events: BoundedSender<SynthesisEvent>,
pub event_receiver: BoundedReceiver<SynthesisEvent>,
}
impl StreamQueues {
pub fn new(capacity: usize) -> Result<Self, EngineError> {
if capacity == 0 {
return Err(EngineError::InvalidConfiguration(
"stream queue capacity must be greater than zero",
));
}
let (pcm, pcm_receiver) = bounded_queue(capacity, StreamKind::Pcm);
let (events, event_receiver) = bounded_queue(capacity, StreamKind::Events);
Ok(Self {
pcm,
pcm_receiver,
events,
event_receiver,
})
}
}
#[derive(Clone)]
pub struct BoundedSender<T> {
kind: StreamKind,
sender: SyncSender<T>,
}
impl<T> BoundedSender<T> {
pub fn send(&self, mut item: T, cancellation: &CancellationToken) -> Result<(), EngineError> {
loop {
cancellation.checkpoint()?;
match self.sender.try_send(item) {
Ok(()) => return Ok(()),
Err(TrySendError::Full(returned)) => {
item = returned;
thread::sleep(BACKPRESSURE_POLL);
}
Err(TrySendError::Disconnected(_)) => {
return Err(EngineError::StreamDisconnected(self.kind));
}
}
}
}
}
pub struct BoundedReceiver<T> {
kind: StreamKind,
receiver: Receiver<T>,
}
impl<T> BoundedReceiver<T> {
pub fn recv_timeout(&self, timeout: Duration) -> Result<T, EngineError> {
match self.receiver.recv_timeout(timeout) {
Ok(item) => Ok(item),
Err(RecvTimeoutError::Timeout) => Err(EngineError::QueueTimeout),
Err(RecvTimeoutError::Disconnected) => Err(EngineError::StreamDisconnected(self.kind)),
}
}
}
fn bounded_queue<T>(capacity: usize, kind: StreamKind) -> (BoundedSender<T>, BoundedReceiver<T>) {
let (sender, receiver) = mpsc::sync_channel(capacity);
(
BoundedSender { kind, sender },
BoundedReceiver { kind, receiver },
)
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum NormalizationMode {
#[default]
Verbatim,
Conservative,
LocaleAware,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct LanguageSpan {
pub range: Range<usize>,
pub language: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PronunciationEntry {
pub language: String,
pub surface: String,
pub spoken: String,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct NormalizationOptions {
pub mode: NormalizationMode,
pub language_spans: Vec<LanguageSpan>,
pub pronunciation_lexicon: Vec<PronunciationEntry>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct NormalizationChange {
pub rule: &'static str,
pub before: String,
pub after: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct NormalizationTrace {
pub mode: NormalizationMode,
pub unicode_version: String,
pub changes: Vec<NormalizationChange>,
}
impl NormalizationTrace {
#[must_use]
pub fn summary(&self) -> NormalizationTraceSummary {
let mut rules = self
.changes
.iter()
.map(|change| change.rule.to_owned())
.collect::<Vec<_>>();
rules.sort_unstable();
rules.dedup();
NormalizationTraceSummary {
mode: self.mode,
unicode_version: self.unicode_version.clone(),
rules,
change_count: self.changes.len(),
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct NormalizationTraceSummary {
pub mode: NormalizationMode,
pub unicode_version: String,
pub rules: Vec<String>,
pub change_count: usize,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PreparedText {
pub token_ids: Vec<u32>,
pub normalization_trace: NormalizationTrace,
}
impl PreparedText {
#[must_use]
pub fn new(token_ids: Vec<u32>, normalization_trace: NormalizationTrace) -> Self {
Self {
token_ids,
normalization_trace,
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TextPreparationError {
message: String,
}
impl TextPreparationError {
#[must_use]
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for TextPreparationError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for TextPreparationError {}
pub trait TextPreparer: Send + Sync {
fn prepare(
&self,
text: &str,
options: &NormalizationOptions,
) -> Result<PreparedText, TextPreparationError>;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CodeFrame {
pub codes: Vec<u32>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GenerationError {
message: String,
}
impl GenerationError {
#[must_use]
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for GenerationError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for GenerationError {}
pub trait FrameGenerator {
fn begin_utterance(&mut self, prepared: &PreparedText) -> Result<(), GenerationError>;
fn next_frame(&mut self) -> Result<Option<CodeFrame>, GenerationError>;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SynthesisRequest {
pub text: String,
pub normalization_options: NormalizationOptions,
pub trace_normalization: bool,
}
impl SynthesisRequest {
#[must_use]
pub fn new(text: impl Into<String>) -> Self {
Self {
text: text.into(),
normalization_options: NormalizationOptions::default(),
trace_normalization: false,
}
}
#[must_use]
pub fn with_normalization_options(
mut self,
normalization_options: NormalizationOptions,
) -> Self {
self.normalization_options = normalization_options;
self
}
#[must_use]
pub const fn with_normalization_trace(mut self, trace_normalization: bool) -> Self {
self.trace_normalization = trace_normalization;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EnrollmentRequest {
pub reference_audio: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SynthesisResult {
pub generated_frames: u64,
pub code_frames: Vec<CodeFrame>,
pub prepared_token_count: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EnrollmentResult {
pub accepted_reference_bytes: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EngineStage {
Synthesis,
Enrollment,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HealthEvent {
BudgetExceeded,
Cancelled,
Violation(health::HealthViolation),
}
impl HealthEvent {
#[must_use]
pub const fn invalidates_output(self) -> bool {
match self {
Self::BudgetExceeded | Self::Cancelled => true,
Self::Violation(violation) => violation.invalidates_output(),
}
}
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::BudgetExceeded => "budget_exceeded",
Self::Cancelled => "cancelled",
Self::Violation(violation) => violation.as_str(),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SynthesisEvent {
Admission { accepted: bool },
ResourceAdmission {
admitted: bool,
predicted_max_frames: u64,
predicted_peak_bytes: u64,
budget_bytes: u64,
},
StageStarted { stage: EngineStage },
StageFinished {
stage: EngineStage,
elapsed: Duration,
},
FrameProgress { frame: u64 },
TextPrepared {
token_count: usize,
normalization: NormalizationTraceSummary,
},
PacketEmitted {
frame_count: u8,
sample_count: usize,
},
Health { event: HealthEvent },
}
pub trait SynthesisObserver: Send + Sync {
fn on_event(&self, event: SynthesisEvent);
}
impl<F> SynthesisObserver for F
where
F: Fn(SynthesisEvent) + Send + Sync,
{
fn on_event(&self, event: SynthesisEvent) {
self(event);
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EngineError {
Busy,
Cancelled,
BudgetExceeded(EngineStage),
StreamDisconnected(StreamKind),
QueueTimeout,
TextPreparation(TextPreparationError),
Generation(GenerationError),
ResourceAdmission(admission::AdmissionRejection),
InvalidConfiguration(&'static str),
Runtime(String),
}
impl fmt::Display for EngineError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Busy => formatter.write_str("another synthesis is already active"),
Self::Cancelled => formatter.write_str("synthesis cancelled"),
Self::BudgetExceeded(stage) => write!(formatter, "{stage:?} stage budget exceeded"),
Self::StreamDisconnected(kind) => write!(formatter, "{kind:?} stream disconnected"),
Self::QueueTimeout => formatter.write_str("bounded queue receive timed out"),
Self::TextPreparation(error) => write!(formatter, "text preparation failed: {error}"),
Self::Generation(error) => write!(formatter, "frame generation failed: {error}"),
Self::ResourceAdmission(rejection) => {
write!(
formatter,
"resource admission refused the request: {rejection}"
)
}
Self::InvalidConfiguration(message) => formatter.write_str(message),
Self::Runtime(message) => write!(formatter, "runtime initialization failed: {message}"),
}
}
}
impl std::error::Error for EngineError {}
pub struct TtsEngine {
runtime: Runtime,
config: EngineConfig,
synthesis_active: AtomicBool,
}
impl TtsEngine {
pub fn new(config: EngineConfig) -> Result<Self, EngineError> {
config.validate()?;
let runtime = RuntimeBuilder::current_thread()
.blocking_threads(1, 1)
.build()
.map_err(|error| EngineError::Runtime(error.to_string()))?;
Ok(Self {
runtime,
config,
synthesis_active: AtomicBool::new(false),
})
}
pub fn from_process_environment() -> Result<Self, EngineError> {
Self::new(process_engine_config())
}
pub fn synthesize<P: TextPreparer + ?Sized>(
&self,
request: SynthesisRequest,
text_preparer: &P,
frame_generator: &mut dyn FrameGenerator,
cancellation: &CancellationToken,
observer: &dyn SynthesisObserver,
) -> Result<SynthesisResult, EngineError> {
let _admission = self.acquire_synthesis_admission(observer)?;
cancellation.checkpoint().inspect_err(|_| {
observer.on_event(SynthesisEvent::Health {
event: HealthEvent::Cancelled,
});
})?;
let prepared = text_preparer
.prepare(&request.text, &request.normalization_options)
.map_err(EngineError::TextPreparation)?;
if request.trace_normalization {
observer.on_event(SynthesisEvent::TextPrepared {
token_count: prepared.token_ids.len(),
normalization: prepared.normalization_trace.summary(),
});
}
let prompt_tokens = prepared.token_ids.len() as u64;
let plan = match self.config.admission.admit(prompt_tokens) {
Ok(plan) => {
observer.on_event(SynthesisEvent::ResourceAdmission {
admitted: true,
predicted_max_frames: plan.predicted_max_frames,
predicted_peak_bytes: plan.predicted_peak_bytes,
budget_bytes: plan.budget_bytes,
});
plan
}
Err(rejection) => {
if let admission::AdmissionRejection::BudgetExceeded { plan } = rejection {
observer.on_event(SynthesisEvent::ResourceAdmission {
admitted: false,
predicted_max_frames: plan.predicted_max_frames,
predicted_peak_bytes: plan.predicted_peak_bytes,
budget_bytes: plan.budget_bytes,
});
}
return Err(EngineError::ResourceAdmission(rejection));
}
};
observer.on_event(SynthesisEvent::StageStarted {
stage: EngineStage::Synthesis,
});
let started = Instant::now();
let startup_budget = self.config.synthesis_stage_budget;
let frame_budget = self.config.synthesis_frame_budget;
let mut code_frames: Vec<CodeFrame> = Vec::new();
frame_generator
.begin_utterance(&prepared)
.map_err(EngineError::Generation)?;
while (code_frames.len() as u64) < plan.predicted_max_frames {
cancellation.checkpoint().inspect_err(|_| {
observer.on_event(SynthesisEvent::Health {
event: HealthEvent::Cancelled,
});
})?;
let deadline = frame_budget
.checked_mul(u32::try_from(code_frames.len()).unwrap_or(u32::MAX))
.and_then(|earned| earned.checked_add(startup_budget))
.unwrap_or(Duration::MAX);
if started.elapsed() > deadline {
observer.on_event(SynthesisEvent::Health {
event: HealthEvent::BudgetExceeded,
});
return Err(EngineError::BudgetExceeded(EngineStage::Synthesis));
}
match frame_generator
.next_frame()
.map_err(EngineError::Generation)?
{
Some(frame) => {
observer.on_event(SynthesisEvent::FrameProgress {
frame: code_frames.len() as u64,
});
code_frames.push(frame);
}
None => break,
}
}
observer.on_event(SynthesisEvent::StageFinished {
stage: EngineStage::Synthesis,
elapsed: started.elapsed(),
});
Ok(SynthesisResult {
generated_frames: code_frames.len() as u64,
code_frames,
prepared_token_count: prepared.token_ids.len(),
})
}
pub fn enroll(
&self,
request: EnrollmentRequest,
cancellation: &CancellationToken,
observer: &dyn SynthesisObserver,
) -> Result<EnrollmentResult, EngineError> {
observer.on_event(SynthesisEvent::Admission { accepted: true });
self.run_stage(
EngineStage::Enrollment,
self.config.enroll_stage_budget,
cancellation,
observer,
|_| Ok(()),
)?;
Ok(EnrollmentResult {
accepted_reference_bytes: request.reference_audio.len(),
})
}
fn acquire_synthesis_admission(
&self,
observer: &dyn SynthesisObserver,
) -> Result<SynthesisAdmission<'_>, EngineError> {
match self.synthesis_active.compare_exchange(
false,
true,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
observer.on_event(SynthesisEvent::Admission { accepted: true });
Ok(SynthesisAdmission { engine: self })
}
Err(_) => {
observer.on_event(SynthesisEvent::Admission { accepted: false });
Err(EngineError::Busy)
}
}
}
fn run_stage<R, F>(
&self,
stage: EngineStage,
budget: Duration,
cancellation: &CancellationToken,
observer: &dyn SynthesisObserver,
work: F,
) -> Result<R, EngineError>
where
R: Send + 'static,
F: FnOnce(CancellationToken) -> Result<R, EngineError> + Send + 'static,
{
cancellation.checkpoint().inspect_err(|_| {
observer.on_event(SynthesisEvent::Health {
event: HealthEvent::Cancelled,
});
})?;
observer.on_event(SynthesisEvent::StageStarted { stage });
let started = Instant::now();
let (sender, receiver) = mpsc::sync_channel(1);
let stage_cancellation = cancellation.clone();
let task_cancellation = cancellation.clone();
let task = self
.runtime
.spawn_blocking(move || {
let result = task_cancellation
.checkpoint()
.and_then(|()| work(task_cancellation));
let _ignored_if_timed_out = sender.send(result);
})
.ok_or_else(|| EngineError::Runtime("blocking pool was not configured".to_owned()))?;
match receiver.recv_timeout(budget) {
Ok(result) => {
let result = result?;
observer.on_event(SynthesisEvent::StageFinished {
stage,
elapsed: started.elapsed(),
});
Ok(result)
}
Err(RecvTimeoutError::Timeout) => {
stage_cancellation.cancel();
task.cancel();
observer.on_event(SynthesisEvent::Health {
event: HealthEvent::BudgetExceeded,
});
Err(EngineError::BudgetExceeded(stage))
}
Err(RecvTimeoutError::Disconnected) => Err(EngineError::Runtime(
"blocking stage disconnected before producing a result".to_owned(),
)),
}
}
}
struct SynthesisAdmission<'a> {
engine: &'a TtsEngine,
}
impl Drop for SynthesisAdmission<'_> {
fn drop(&mut self) {
self.engine.synthesis_active.store(false, Ordering::Release);
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
#[derive(Default)]
struct RecordingObserver {
events: Mutex<Vec<SynthesisEvent>>,
}
impl RecordingObserver {
fn events(&self) -> Vec<SynthesisEvent> {
self.events
.lock()
.expect("test observer lock poisoned")
.clone()
}
}
impl SynthesisObserver for RecordingObserver {
fn on_event(&self, event: SynthesisEvent) {
self.events
.lock()
.expect("test observer lock poisoned")
.push(event);
}
}
fn engine_with_budget(budget: Duration) -> TtsEngine {
TtsEngine::new(EngineConfig {
synthesis_stage_budget: budget,
..EngineConfig::default()
})
.expect("test engine builds")
}
fn engine_with_frame_budget(startup: Duration, per_frame: Duration) -> TtsEngine {
TtsEngine::new(EngineConfig {
synthesis_stage_budget: startup,
synthesis_frame_budget: per_frame,
..EngineConfig::default()
})
.expect("test engine builds")
}
struct PacedFrameGenerator {
remaining: usize,
per_frame: Duration,
stall: Option<Duration>,
began: bool,
}
impl FrameGenerator for PacedFrameGenerator {
fn begin_utterance(&mut self, _prepared: &PreparedText) -> Result<(), GenerationError> {
self.began = true;
Ok(())
}
fn next_frame(&mut self) -> Result<Option<CodeFrame>, GenerationError> {
assert!(self.began, "next_frame before begin_utterance");
if self.remaining == 0 {
let Some(stall) = self.stall else {
return Ok(None);
};
thread::sleep(stall);
return Ok(Some(CodeFrame { codes: vec![0; 16] }));
}
self.remaining -= 1;
thread::sleep(self.per_frame);
Ok(Some(CodeFrame { codes: vec![0; 16] }))
}
}
#[test]
fn steady_progress_past_the_startup_grace_is_not_refused_for_being_long() {
let engine = engine_with_frame_budget(Duration::from_millis(50), Duration::from_millis(30));
let observer = RecordingObserver::default();
let mut generator = PacedFrameGenerator {
remaining: 10,
per_frame: Duration::from_millis(20),
stall: None,
began: false,
};
let result = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut generator,
&CancellationToken::new(),
&observer,
)
.expect("a steadily-progressing run must not be refused");
assert_eq!(
result.generated_frames, 10,
"all ten frames must survive; a flat 50 ms ceiling would have cut this at ~2"
);
}
#[test]
fn a_generator_that_stops_progressing_is_still_caught_within_its_earned_deadline() {
let engine = engine_with_frame_budget(Duration::from_millis(50), Duration::from_millis(30));
let observer = RecordingObserver::default();
let mut generator = PacedFrameGenerator {
remaining: 3,
per_frame: Duration::from_millis(1),
stall: Some(Duration::from_millis(400)),
began: false,
};
let started = Instant::now();
let error = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut generator,
&CancellationToken::new(),
&observer,
)
.expect_err("a stalled generator must still be refused");
let elapsed = started.elapsed();
assert_eq!(error, EngineError::BudgetExceeded(EngineStage::Synthesis));
assert!(
observer.events().contains(&SynthesisEvent::Health {
event: HealthEvent::BudgetExceeded,
}),
"the stall must be reported on the health channel, not only as a return value"
);
assert!(
elapsed < Duration::from_millis(2000),
"stall detection took {elapsed:?}; the deadline is not supposed to keep growing while \
no frames are produced"
);
}
#[test]
fn a_zero_frame_budget_is_rejected_rather_than_collapsing_to_a_flat_deadline() {
let built = TtsEngine::new(EngineConfig {
synthesis_frame_budget: Duration::ZERO,
..EngineConfig::default()
});
assert!(
matches!(built, Err(EngineError::InvalidConfiguration(_))),
"a zero per-frame budget must be refused, not accepted as a flat deadline"
);
}
#[test]
fn an_unoptimized_build_is_granted_a_larger_synthesis_budget() {
let config = EngineConfig::default();
let expected = u32::from(cfg!(debug_assertions)) * (DEBUG_BUILD_SLOWDOWN - 1) + 1;
assert_eq!(
config.synthesis_frame_budget,
DEFAULT_SYNTHESIS_FRAME_BUDGET * expected
);
assert_eq!(
config.synthesis_stage_budget,
DEFAULT_SYNTHESIS_BUDGET * expected
);
assert_eq!(config.enroll_stage_budget, DEFAULT_ENROLL_BUDGET);
}
struct ScriptedFrameGenerator {
remaining: usize,
began: bool,
endless: bool,
polls: usize,
}
impl ScriptedFrameGenerator {
fn emitting(frames: usize) -> Self {
Self {
remaining: frames,
began: false,
endless: false,
polls: 0,
}
}
fn endless() -> Self {
Self {
remaining: 0,
began: false,
endless: true,
polls: 0,
}
}
}
impl FrameGenerator for ScriptedFrameGenerator {
fn begin_utterance(&mut self, _prepared: &PreparedText) -> Result<(), GenerationError> {
self.began = true;
Ok(())
}
fn next_frame(&mut self) -> Result<Option<CodeFrame>, GenerationError> {
assert!(self.began, "next_frame before begin_utterance");
self.polls += 1;
if self.endless {
return Ok(Some(CodeFrame { codes: vec![0; 16] }));
}
if self.remaining == 0 {
return Ok(None);
}
self.remaining -= 1;
Ok(Some(CodeFrame { codes: vec![0; 16] }))
}
}
fn engine_with_frame_ceiling(max_new_tokens: u64) -> TtsEngine {
TtsEngine::new(EngineConfig {
synthesis_stage_budget: Duration::from_secs(5),
admission: admission::AdmissionPolicy {
max_new_tokens,
..admission::AdmissionPolicy::default()
},
..EngineConfig::default()
})
.expect("test engine builds")
}
fn admitted_ceiling(observer: &RecordingObserver) -> u64 {
observer
.events()
.into_iter()
.find_map(|event| match event {
SynthesisEvent::ResourceAdmission {
admitted: true,
predicted_max_frames,
..
} => Some(predicted_max_frames),
_ => None,
})
.expect("an admitted request reports its frame ceiling")
}
struct TestTextPreparer;
impl TextPreparer for TestTextPreparer {
fn prepare(
&self,
_text: &str,
options: &NormalizationOptions,
) -> Result<PreparedText, TextPreparationError> {
Ok(PreparedText::new(
vec![7, 11],
NormalizationTrace {
mode: options.mode,
unicode_version: "15.1.0".to_owned(),
changes: vec![NormalizationChange {
rule: "unicode_nfc",
before: "caller-owned secret".to_owned(),
after: "caller-owned secret".to_owned(),
}],
},
))
}
}
#[test]
fn the_decode_loop_stops_on_the_generators_eos_and_polls_exactly_once_past_it() {
let engine = engine_with_frame_ceiling(64);
let observer = RecordingObserver::default();
let mut generator = ScriptedFrameGenerator::emitting(3);
let result = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut generator,
&CancellationToken::new(),
&observer,
)
.expect("scripted pipeline succeeds");
let ceiling = admitted_ceiling(&observer);
assert!(
ceiling > 3,
"ceiling {ceiling} must exceed the 3 emitted frames, or the stop is ambiguous"
);
assert_eq!(result.generated_frames, 3, "EOS bounds the utterance");
assert_eq!(result.code_frames.len(), 3);
assert_eq!(
generator.polls, 4,
"the loop must poll once past the last frame to observe EOS, and then stop"
);
}
#[test]
fn a_generator_that_never_stops_is_truncated_exactly_at_the_admitted_ceiling() {
let engine = engine_with_frame_ceiling(5);
let observer = RecordingObserver::default();
let mut generator = ScriptedFrameGenerator::endless();
let result = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut generator,
&CancellationToken::new(),
&observer,
)
.expect("a ceiling-bound utterance still completes");
let ceiling = admitted_ceiling(&observer);
assert_eq!(
ceiling, 5,
"the policy's max_new_tokens is the ceiling here"
);
assert_eq!(
result.generated_frames, ceiling,
"an endless generator must be cut at the ceiling, not one frame either side"
);
assert_eq!(result.code_frames.len() as u64, ceiling);
assert_eq!(
generator.polls as u64, ceiling,
"once the ceiling is reached the loop must stop asking, not poll a discarded frame"
);
}
#[test]
fn eos_landing_exactly_on_the_ceiling_yields_the_ceiling_frames() {
let engine = engine_with_frame_ceiling(4);
let observer = RecordingObserver::default();
let mut generator = ScriptedFrameGenerator::emitting(4);
let result = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut generator,
&CancellationToken::new(),
&observer,
)
.expect("scripted pipeline succeeds");
assert_eq!(admitted_ceiling(&observer), 4);
assert_eq!(result.generated_frames, 4);
assert_eq!(
generator.polls, 4,
"the ceiling is reached first, so the generator is never asked for a fifth frame"
);
}
#[test]
fn the_decode_loop_drives_the_generator_and_reports_every_frame() {
let engine = engine_with_budget(Duration::from_secs(1));
let cancellation = CancellationToken::new();
let observer = RecordingObserver::default();
let mut generator = ScriptedFrameGenerator::emitting(2);
let result = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut generator,
&cancellation,
&observer,
)
.expect("scripted pipeline succeeds");
assert_eq!(result.generated_frames, 2);
assert_eq!(result.code_frames.len(), 2);
assert!(
result
.code_frames
.iter()
.all(|frame| frame.codes.len() == 16)
);
assert_eq!(result.prepared_token_count, 2);
let events = observer.events();
assert!(
matches!(
events.as_slice(),
[
SynthesisEvent::Admission { accepted: true },
SynthesisEvent::ResourceAdmission { admitted: true, .. },
SynthesisEvent::StageStarted {
stage: EngineStage::Synthesis,
},
SynthesisEvent::FrameProgress { frame: 0 },
SynthesisEvent::FrameProgress { frame: 1 },
SynthesisEvent::StageFinished {
stage: EngineStage::Synthesis,
..
},
]
),
"unexpected event sequence: {events:?}"
);
}
#[test]
fn an_unaffordable_request_is_refused_before_any_stage_runs() {
let mut config = EngineConfig {
synthesis_stage_budget: Duration::from_secs(1),
..EngineConfig::default()
};
config.admission.budget_bytes = 1;
let engine = TtsEngine::new(config).expect("engine builds");
let cancellation = CancellationToken::new();
let observer = RecordingObserver::default();
let error = engine
.synthesize(
SynthesisRequest::new(""),
&TestTextPreparer,
&mut ScriptedFrameGenerator::emitting(0),
&cancellation,
&observer,
)
.expect_err("an unaffordable request must be refused");
assert!(
matches!(error, EngineError::ResourceAdmission(_)),
"got {error}"
);
let events = observer.events();
assert!(
!events.iter().any(|event| matches!(
event,
SynthesisEvent::StageStarted { .. }
| SynthesisEvent::StageFinished { .. }
| SynthesisEvent::FrameProgress { .. }
)),
"a refused request must not start any stage; got {events:?}"
);
assert!(
events.iter().any(|event| matches!(
event,
SynthesisEvent::ResourceAdmission {
admitted: false,
..
}
)),
"a capacity refusal must appear in the event stream: {events:?}"
);
}
#[test]
fn the_admission_policy_is_configurable_and_defaults_are_documented() {
let config = EngineConfig::default();
assert_eq!(
config.admission.budget_bytes,
admission::DEFAULT_BUDGET_BYTES
);
assert_eq!(
config.admission.max_new_tokens,
admission::DEFAULT_MAX_NEW_TOKENS
);
let plan = config
.admission
.admit(512)
.expect("the documented default must admit its own worked case");
assert_eq!(
plan.predicted_max_frames,
512 * admission::HEURISTIC_FRAMES_PER_PROMPT_TOKEN
+ admission::HEURISTIC_FRAME_HEADROOM
);
assert!(plan.fits());
let mut explicit = config.admission;
explicit.heuristic_eos_backstop = false;
let plan = explicit
.admit(512)
.expect("the documented explicit-cap case must admit");
assert_eq!(plan.predicted_max_frames, admission::DEFAULT_MAX_NEW_TOKENS);
assert!(plan.fits());
}
#[test]
fn cancellation_is_observed_before_the_cpu_stage_starts() {
let engine = engine_with_budget(Duration::from_secs(1));
let cancellation = CancellationToken::new();
cancellation.cancel();
let observer = RecordingObserver::default();
let error = engine
.synthesize(
SynthesisRequest::new("cancelled"),
&TestTextPreparer,
&mut ScriptedFrameGenerator::emitting(0),
&cancellation,
&observer,
)
.expect_err("cancelled request must not run");
assert_eq!(error, EngineError::Cancelled);
assert_eq!(
observer.events(),
vec![
SynthesisEvent::Admission { accepted: true },
SynthesisEvent::Health {
event: HealthEvent::Cancelled,
},
]
);
}
#[test]
fn stage_budget_cancels_cooperative_cpu_work() {
let engine = engine_with_budget(Duration::from_millis(5));
let cancellation = CancellationToken::new();
let observer = RecordingObserver::default();
let error = engine
.run_stage(
EngineStage::Synthesis,
Duration::from_millis(5),
&cancellation,
&observer,
|token| -> Result<(), EngineError> {
loop {
token.checkpoint()?;
thread::sleep(Duration::from_millis(1));
}
},
)
.expect_err("long stage must time out");
assert_eq!(error, EngineError::BudgetExceeded(EngineStage::Synthesis));
assert!(cancellation.is_cancelled());
assert!(observer.events().contains(&SynthesisEvent::Health {
event: HealthEvent::BudgetExceeded,
}));
}
#[test]
fn pcm_and_events_have_independent_bounded_queues() {
let queues = StreamQueues::new(1).expect("queue config is valid");
let cancellation = CancellationToken::new();
queues
.events
.send(SynthesisEvent::Admission { accepted: true }, &cancellation)
.expect("event queue accepts first event");
queues
.pcm
.send(
PcmPacket {
frame_count: 1,
samples: vec![1, -1],
},
&cancellation,
)
.expect("full event queue cannot block PCM queue");
assert_eq!(
queues
.pcm_receiver
.recv_timeout(Duration::from_millis(10))
.expect("PCM arrives"),
PcmPacket {
frame_count: 1,
samples: vec![1, -1],
}
);
}
#[test]
fn explicit_normalization_trace_is_text_free() {
let engine = engine_with_budget(Duration::from_secs(1));
let observer = RecordingObserver::default();
let request = SynthesisRequest::new("caller-owned secret")
.with_normalization_options(NormalizationOptions {
mode: NormalizationMode::LocaleAware,
..NormalizationOptions::default()
})
.with_normalization_trace(true);
engine
.synthesize(
request,
&TestTextPreparer,
&mut ScriptedFrameGenerator::emitting(0),
&CancellationToken::new(),
&observer,
)
.expect("explicit trace request succeeds");
let trace = observer
.events()
.into_iter()
.find_map(|event| match event {
SynthesisEvent::TextPrepared {
token_count,
normalization,
} => Some((token_count, normalization)),
_ => None,
})
.expect("explicit request emits a trace summary");
assert_eq!(trace.0, 2);
assert_eq!(trace.1.mode, NormalizationMode::LocaleAware);
assert_eq!(trace.1.unicode_version, "15.1.0");
assert_eq!(trace.1.rules, vec!["unicode_nfc"]);
assert_eq!(trace.1.change_count, 1);
assert!(
!format!("{:?}", trace.1).contains("caller-owned secret"),
"observer trace summaries must never contain sensitive before/after text"
);
}
#[test]
fn a_health_violation_reaches_the_caller_through_the_observer() {
let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = Arc::clone(&seen);
let observer = move |event: SynthesisEvent| {
if let SynthesisEvent::Health { event } = event {
sink.lock().expect("observer lock").push(event);
}
};
let violation = health::HealthViolation::OutputSilent {
silent_millis: 1_500,
};
observer(SynthesisEvent::Health {
event: HealthEvent::Violation(violation),
});
let demotion = health::HealthViolation::KernelDemoted {
from: health::KernelTier::Optimized("i8mm"),
to: health::KernelTier::Scalar,
};
observer(SynthesisEvent::Health {
event: HealthEvent::Violation(demotion),
});
let events = seen.lock().expect("observer lock").clone();
assert_eq!(events.len(), 2);
assert_eq!(events[0], HealthEvent::Violation(violation));
assert_eq!(events[0].as_str(), "output_silent");
assert!(events[0].invalidates_output());
assert!(!events[1].invalidates_output());
assert_eq!(events[1].as_str(), "kernel_demoted");
}
#[test]
fn many_utterances_without_deadlock_watchdog() {
let (done_sender, done_receiver) = mpsc::sync_channel(1);
let worker = thread::spawn(move || {
let engine = engine_with_budget(Duration::from_secs(1));
let observer = RecordingObserver::default();
for _ in 0..64 {
engine
.synthesize(
SynthesisRequest::new("watchdog"),
&TestTextPreparer,
&mut ScriptedFrameGenerator::emitting(1),
&CancellationToken::new(),
&observer,
)
.expect("empty utterance succeeds");
}
done_sender
.send(())
.expect("watchdog completion receiver lives");
});
done_receiver
.recv_timeout(Duration::from_secs(2))
.expect("many utterances watchdog expired");
worker.join().expect("watchdog worker does not panic");
}
}