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

//! Bench harness for converse_server. Drives N sequential turns through
//! a single WebSocket and reports per-phase latency stats:
//!
//!   - audio_send_ms     (client streaming PCM in until EOT)
//!   - transcript_ms     (server time from EOT → "transcript" event)
//!   - first_audio_ms    (server time from "transcript" event → first audio chunk)
//!   - turn_total_ms     (audio_send + transcript + LLM+TTS to "done")
//!
//! At end, fetches /metrics from the server and prints the summary.
//!
//! Usage:
//! ```bash
//! # Terminal 1: server with mock LLM
//! cargo run -p rtx-csm --release --features metal --example converse_server -- \
//!     --bind 127.0.0.1:18096 --mock-llm
//!
//! # Terminal 2: bench
//! cargo run -p rtx-csm --release --example converse_server_bench -- \
//!     --base http://127.0.0.1:18096 --turns 3
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use futures_util::{SinkExt, StreamExt};
use rtx_csm::audio_io;
use std::path::PathBuf;
use std::time::{Duration, Instant};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::protocol::Message;

const PCM_RATE: u32 = 24_000;
const FRAME_SAMPLES: usize = 4800; // 200 ms

#[derive(Debug, Parser)]
struct Cli {
    /// Server base URL (http://...).
    #[arg(long, default_value = "http://127.0.0.1:18090")]
    base: String,
    /// Number of sequential turns.
    #[arg(long, default_value_t = 3)]
    turns: usize,
    /// User-turn audio (any rate; resampled to 24 kHz).
    #[arg(long = "in", default_value = "/tmp/asr_test.flac")]
    input: PathBuf,
    /// Bearer token. Reads RTX_AUTH_TOKEN if not set.
    #[arg(long)]
    auth_token: Option<String>,
    /// Send audio at real-time pace (sleep between frames to match
    /// audio playback rate). Required to measure the parallel-STT win;
    /// without it the client dumps all audio at once and the server
    /// processes serially regardless.
    #[arg(long)]
    realtime: bool,
}

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

impl Stats {
    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).min((n - 1) as f64) as usize];
        let mean = s.iter().sum::<f64>() / n as f64;
        let min = s[0];
        let max = s[n - 1];
        println!(
            "  {} (n={n}): mean={mean:.0}ms p50={p50:.0}ms p95={p95:.0}ms min={min:.0}ms max={max:.0}ms",
            self.label
        );
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    let samples = audio_io::load_mono_at_rate(&cli.input, PCM_RATE).context("load user audio")?;
    let n_samples = samples.len();
    println!(
        "loaded {}: {} samples ({:.2}s @ {} Hz)",
        cli.input.display(),
        n_samples,
        n_samples as f32 / PCM_RATE as f32,
        PCM_RATE
    );

    let ws_url = cli
        .base
        .replace("http://", "ws://")
        .replace("https://", "wss://");
    let ws_url = format!("{}/v1/converse", ws_url.trim_end_matches('/'));
    println!("== converse_server bench ==");
    println!("base={} ws={ws_url} turns={}", cli.base, cli.turns);

    let mut req = (&ws_url).into_client_request().context("parse url")?;
    if let Some(token) = cli
        .auth_token
        .clone()
        .or_else(|| std::env::var("RTX_AUTH_TOKEN").ok())
    {
        req.headers_mut().insert(
            "Authorization",
            format!("Bearer {token}").parse().context("auth header")?,
        );
    }
    let (mut ws, _resp) = tokio_tungstenite::connect_async(req)
        .await
        .with_context(|| format!("connect {ws_url}"))?;
    println!("connected");

    let mut send_lat = Stats::new("audio_send");
    let mut tx_lat = Stats::new("transcript_ms");
    let mut ttfa_lat = Stats::new("first_audio_ms");
    let mut total_lat = Stats::new("turn_total_ms");

    for i in 0..cli.turns {
        let turn_t = Instant::now();
        // Send audio frames + EOT.
        let send_t = Instant::now();
        let frame_duration = if cli.realtime {
            Some(Duration::from_secs_f64(
                FRAME_SAMPLES as f64 / PCM_RATE as f64,
            ))
        } else {
            None
        };
        let frame_pace_start = Instant::now();
        for (idx, chunk) in samples.chunks(FRAME_SAMPLES).enumerate() {
            let mut buf = Vec::with_capacity(chunk.len() * 2);
            for &s in chunk {
                let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16;
                buf.extend_from_slice(&v.to_le_bytes());
            }
            // For realtime pacing, sleep until this frame's playback time.
            if let Some(d) = frame_duration {
                let target = frame_pace_start + d * (idx as u32 + 1);
                let now = Instant::now();
                if target > now {
                    tokio::time::sleep(target - now).await;
                }
            }
            ws.send(Message::Binary(buf.into())).await?;
        }
        ws.send(Message::Text("EOT".into())).await?;
        let send_ms = send_t.elapsed();

        // Wait for transcript event, first audio chunk, done event.
        let mut transcript_t: Option<Instant> = None;
        let mut first_audio_t: Option<Instant> = None;
        loop {
            match ws.next().await {
                Some(Ok(Message::Binary(_))) => {
                    if first_audio_t.is_none() {
                        first_audio_t = Some(Instant::now());
                    }
                }
                Some(Ok(Message::Text(t))) => {
                    let v: serde_json::Value = serde_json::from_str(&t)?;
                    let event = v.get("event").and_then(|x| x.as_str()).unwrap_or("");
                    match event {
                        "transcript" => transcript_t = Some(Instant::now()),
                        "done" => break,
                        "error" => {
                            let msg = v.get("msg").and_then(|x| x.as_str()).unwrap_or("(unknown)");
                            anyhow::bail!("server error: {msg}");
                        }
                        _ => {}
                    }
                }
                Some(Ok(_)) => {}
                Some(Err(e)) => return Err(e.into()),
                None => break,
            }
        }
        let total_ms = turn_t.elapsed();
        send_lat.add(send_ms);
        if let Some(tt) = transcript_t {
            // tx_ms = transcript event time relative to send completion.
            tx_lat.add(tt.duration_since(send_t + send_ms));
            if let Some(fa) = first_audio_t {
                ttfa_lat.add(fa.duration_since(tt));
            }
        }
        total_lat.add(total_ms);
        println!(
            "  turn {}: audio_send={:.0}ms total={:.0}ms",
            i + 1,
            send_ms.as_secs_f64() * 1000.0,
            total_ms.as_secs_f64() * 1000.0
        );
    }

    let _ = ws.close(None).await;
    println!("\n--- per-phase stats ---");
    send_lat.report();
    tx_lat.report();
    ttfa_lat.report();
    total_lat.report();

    // Pull /metrics
    let metrics_url = format!("{}/metrics", cli.base.trim_end_matches('/'));
    let body = reqwest::Client::new()
        .get(&metrics_url)
        .send()
        .await?
        .text()
        .await?;

    // Pretty-print just the per-phase server-side gauges. These are
    // averages over the lifetime of the server (i.e. all turns we just
    // ran, since the server has been single-use). They complement the
    // client-side stats above by including time the client can't see
    // (mutex contention, internal mpsc, etc.).
    let phase_keys = [
        ("rtx_csm_recv_phase_ms_avg", "recv_phase"),
        ("rtx_csm_llm_to_first_audio_ms_avg", "llm_to_first_audio"),
        ("rtx_csm_conv_total_ms_avg", "conv_total"),
        ("rtx_csm_total_turn_ms_avg", "total_turn"),
        ("rtx_csm_stt_latency_ms_avg", "stt_post"),
        ("rtx_csm_tts_latency_ms_avg", "tts_per_utterance"),
        ("rtx_csm_e2e_first_audio_ms_avg", "e2e_first_audio"),
    ];
    println!("\n--- server-side phase averages (across all turns) ---");
    for (key, label) in phase_keys.iter() {
        if let Some(v) = parse_gauge(&body, key) {
            println!("  {label:<22}{v:>6}ms");
        }
    }

    // If the user wants the raw text we still emit it, but at the
    // bottom — most of the time the parsed view is what's wanted.
    println!("\n--- raw /metrics ---\n{body}");

    Ok(())
}

/// Extract a single Prometheus gauge value (`name <int>` or `name <float>`).
/// Comments (`# ...`) and unrelated lines are skipped. Returns `None` if
/// the gauge isn't present (e.g., zero turns ran).
fn parse_gauge(body: &str, name: &str) -> Option<u64> {
    for line in body.lines() {
        let line = line.trim_start();
        if line.starts_with('#') {
            continue;
        }
        let mut parts = line.split_whitespace();
        if parts.next()? != name {
            continue;
        }
        let v = parts.next()?;
        // Gauges are written as integers in the server but accept floats.
        if let Ok(i) = v.parse::<u64>() {
            return Some(i);
        }
        if let Ok(f) = v.parse::<f64>() {
            return Some(f as u64);
        }
    }
    None
}