rustytorch/crates/models/rtx-csm/examples/tts_server_bench.rs

//! End-to-end bench for the tts_server endpoints.
//!
//! Runs a small request mix (1 sequential + N concurrent) and prints
//! per-endpoint p50/p95/throughput stats. The server should already be
//! running with `--audioseal-*` and `--wavlm-sv` flags so all endpoints
//! are enabled.
//!
//! Usage:
//! ```
//! # Terminal 1:
//! cargo run -p rtx-csm --release --features metal --example tts_server -- \
//!     --bind 127.0.0.1:18080 \
//!     --audioseal-generator /tmp/audioseal_generator.safetensors \
//!     --audioseal-detector  /tmp/audioseal_detector.safetensors \
//!     --wavlm-sv /tmp/wavlm_sv.safetensors
//!
//! # Terminal 2:
//! cargo run -p rtx-csm --release --example tts_server_bench -- \
//!     --base http://127.0.0.1:18080 --concurrent 4 --tts-runs 6
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};

#[derive(Debug, Parser)]
struct Cli {
    /// Server base URL (no trailing slash).
    #[arg(long, default_value = "http://127.0.0.1:18080")]
    base: String,
    /// Number of concurrent requests for the concurrency test.
    #[arg(long, default_value_t = 4)]
    concurrent: usize,
    /// Number of sequential TTS requests to run for latency stats.
    #[arg(long, default_value_t = 6)]
    tts_runs: usize,
    /// max_audio_ms for each TTS request.
    #[arg(long, default_value_t = 3000)]
    tts_ms: u32,
    /// Optional WAV used for /v1/detect, /v1/speaker_embed, /v1/speaker_compare.
    /// Defaults to /tmp/srv_bench_input.wav (auto-generated below if missing).
    #[arg(long, default_value = "/tmp/srv_bench_input.wav")]
    audio_for_aux: PathBuf,
}

#[derive(Debug, Default)]
struct LatencyStats {
    label: String,
    samples: Vec<f64>,
}

impl LatencyStats {
    fn new(label: &str) -> Self {
        Self {
            label: label.to_string(),
            samples: Vec::new(),
        }
    }
    fn add(&mut self, d: Duration) {
        self.samples.push(d.as_secs_f64() * 1000.0);
    }
    fn report(&self) {
        if self.samples.is_empty() {
            println!("  {}: no samples", self.label);
            return;
        }
        let mut s = self.samples.clone();
        s.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let n = s.len();
        let p50 = s[n / 2];
        let p95 = s[((n as f64) * 0.95) as usize];
        let mean = s.iter().sum::<f64>() / n as f64;
        let min = s[0];
        let max = s[n - 1];
        println!(
            "  {} (n={}): mean={:.1}ms p50={:.1}ms p95={:.1}ms min={:.1}ms max={:.1}ms",
            self.label, n, mean, p50, p95, min, max
        );
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(120))
        .build()?;

    println!("== rtx-csm tts_server bench ==");
    println!(
        "base={} concurrent={} tts_runs={}",
        cli.base, cli.concurrent, cli.tts_runs
    );

    // Health probe.
    let t = Instant::now();
    let body = client
        .get(format!("{}/health", cli.base))
        .send()
        .await
        .context("health: connection failed (is the server running?)")?
        .text()
        .await?;
    println!(
        "\n[health] {:.1}ms -> {body:?}",
        t.elapsed().as_secs_f64() * 1000.0
    );

    // -- Sequential /v1/tts latency stats ---------------------------------
    println!("\n--- sequential /v1/tts (max_audio_ms={}) ---", cli.tts_ms);
    let mut tts_lat = LatencyStats::new("/v1/tts");
    let mut last_wav: Vec<u8> = Vec::new();
    for i in 0..cli.tts_runs {
        let t = Instant::now();
        let res = client
            .post(format!("{}/v1/tts", cli.base))
            .json(&serde_json::json!({
                "text": format!("Bench request number {}.", i + 1),
                "speaker": 0,
                "max_audio_ms": cli.tts_ms,
                "seed": 42 + i as u64,
            }))
            .send()
            .await?;
        let status = res.status();
        let bytes = res.bytes().await?;
        let dt = t.elapsed();
        tts_lat.add(dt);
        println!(
            "  run {}: {} bytes={} {:.0}ms",
            i + 1,
            status,
            bytes.len(),
            dt.as_secs_f64() * 1000.0
        );
        if i == 0 {
            last_wav = bytes.to_vec();
            std::fs::write(&cli.audio_for_aux, &last_wav).ok();
        }
    }
    tts_lat.report();

    // -- /v1/detect latency ----------------------------------------------
    if !last_wav.is_empty() {
        println!("\n--- /v1/detect ---");
        let mut det_lat = LatencyStats::new("/v1/detect");
        for _ in 0..cli.tts_runs {
            let part = reqwest::multipart::Part::bytes(last_wav.clone())
                .file_name("audio.wav")
                .mime_str("audio/wav")?;
            let form = reqwest::multipart::Form::new().part("audio", part);
            let t = Instant::now();
            let body = client
                .post(format!("{}/v1/detect", cli.base))
                .multipart(form)
                .send()
                .await?
                .text()
                .await?;
            det_lat.add(t.elapsed());
            // Show one sample only.
            if det_lat.samples.len() == 1 {
                println!("  sample response: {body}");
            }
        }
        det_lat.report();

        // -- /v1/speaker_embed latency -----------------------------------
        println!("\n--- /v1/speaker_embed ---");
        let mut emb_lat = LatencyStats::new("/v1/speaker_embed");
        for _ in 0..cli.tts_runs {
            let part = reqwest::multipart::Part::bytes(last_wav.clone())
                .file_name("audio.wav")
                .mime_str("audio/wav")?;
            let form = reqwest::multipart::Form::new().part("audio", part);
            let t = Instant::now();
            let _ = client
                .post(format!("{}/v1/speaker_embed", cli.base))
                .multipart(form)
                .send()
                .await?
                .bytes()
                .await?;
            emb_lat.add(t.elapsed());
        }
        emb_lat.report();

        // -- /v1/speaker_compare latency ---------------------------------
        println!("\n--- /v1/speaker_compare (a == b) ---");
        let mut cmp_lat = LatencyStats::new("/v1/speaker_compare");
        for _ in 0..cli.tts_runs {
            let pa = reqwest::multipart::Part::bytes(last_wav.clone())
                .file_name("a.wav")
                .mime_str("audio/wav")?;
            let pb = reqwest::multipart::Part::bytes(last_wav.clone())
                .file_name("b.wav")
                .mime_str("audio/wav")?;
            let form = reqwest::multipart::Form::new().part("a", pa).part("b", pb);
            let t = Instant::now();
            let body = client
                .post(format!("{}/v1/speaker_compare", cli.base))
                .multipart(form)
                .send()
                .await?
                .text()
                .await?;
            cmp_lat.add(t.elapsed());
            if cmp_lat.samples.len() == 1 {
                println!("  sample response: {body}");
            }
        }
        cmp_lat.report();
    }

    // -- Concurrent /v1/tts ---------------------------------------------
    println!("\n--- concurrent /v1/tts (n={}) ---", cli.concurrent);
    let client = Arc::new(client);
    let base = Arc::new(cli.base.clone());
    let t_total = Instant::now();
    let mut handles = Vec::with_capacity(cli.concurrent);
    for i in 0..cli.concurrent {
        let client = client.clone();
        let base = base.clone();
        let tts_ms = cli.tts_ms;
        handles.push(tokio::spawn(async move {
            let t = Instant::now();
            let res = client
                .post(format!("{}/v1/tts", base))
                .json(&serde_json::json!({
                    "text": format!("Concurrent request {}.", i + 1),
                    "speaker": 0,
                    "max_audio_ms": tts_ms,
                    "seed": 100 + i as u64,
                }))
                .send()
                .await
                .map_err(|e| format!("send: {e}"))?;
            let status = res.status();
            let bytes = res.bytes().await.map_err(|e| format!("body: {e}"))?.len();
            Ok::<_, String>((i, t.elapsed(), status, bytes))
        }));
    }
    let mut conc_lat = LatencyStats::new("concurrent /v1/tts");
    for h in handles {
        match h.await? {
            Ok((i, dt, status, bytes)) => {
                conc_lat.add(dt);
                println!(
                    "  worker {}: {} bytes={} {:.0}ms",
                    i,
                    status,
                    bytes,
                    dt.as_secs_f64() * 1000.0
                );
            }
            Err(e) => println!("  worker error: {e}"),
        }
    }
    let total = t_total.elapsed();
    conc_lat.report();
    println!(
        "  wall-clock total: {:.1}s  (effective serial = {:.1}s)",
        total.as_secs_f64(),
        conc_lat.samples.iter().sum::<f64>() / 1000.0,
    );
    println!(
        "  serialization factor = {:.2}x  (1.0 = perfectly parallel; >1.0 = Mutex-serialized)",
        (conc_lat.samples.iter().sum::<f64>() / 1000.0) / total.as_secs_f64()
    );

    Ok(())
}