inferencelayer 0.2.10

//! Concurrent-decode load generator — the Rust replacement for the Python harness.
//!
//! The throughput rows in this campaign were driven by a CPython script running one thread per
//! stream, each parsing a server-sent event per generated token. At M=64 that is ~12k GIL-bound
//! event parses a second, which is a measurement instrument capable of capping the thing it
//! measures. It did not, in fact, cap it — a non-streaming variant of the same script produced
//! 11604 tok/s against 11681 streamed — but "I checked and it was fine" is a weaker claim than
//! not having the risk at all, and every number that goes in a report against vLLM and SGLang
//! should come off a client that cannot be the bottleneck.
//!
//! Speaks plain HTTP/1.1 over `std::net::TcpStream` (no client crate, no async runtime) to the
//! OpenAI `/v1/completions` endpoint, so the SAME binary drives ours, vLLM and SGLang. Chunked
//! transfer-encoding is decoded properly rather than by scanning for `data:` in the raw stream —
//! generated text can contain that substring, and a miscount would be invisible.
//!
//! usage: loadgen <port> <ngen> <M> <ids-file> [--nostream] [--host H]
//!
//! Reports the same quantity as the Python harness: total tokens over the wall clock of the
//! concurrent burst, plus per-stream TTFT so the prefill and decode axes come off one run.
use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpStream;
use std::time::Instant;

fn arg_flag(a: &[String], name: &str) -> bool { a.iter().any(|s| s == name) }
fn arg_val(a: &[String], name: &str) -> Option<String> {
    a.iter().position(|s| s == name).and_then(|i| a.get(i + 1)).cloned()
}

/// Read one HTTP response: headers, then a body that is either chunked or Content-Length'd.
/// `on_event` fires per SSE `data:` payload (streaming); the full body is returned otherwise.
fn read_response<F: FnMut(&str)>(r: &mut BufReader<TcpStream>, mut on_event: F) -> std::io::Result<Vec<u8>> {
    let mut chunked = false;
    let mut clen: Option<usize> = None;
    loop {
        let mut line = String::new();
        if r.read_line(&mut line)? == 0 { return Ok(Vec::new()) }
        let t = line.trim_end();
        if t.is_empty() { break }
        let lower = t.to_ascii_lowercase();
        if let Some(v) = lower.strip_prefix("transfer-encoding:") {
            chunked = v.trim() == "chunked";
        } else if let Some(v) = lower.strip_prefix("content-length:") {
            clen = v.trim().parse().ok();
        }
    }

    let mut body = Vec::new();
    let mut pending = Vec::new();   // partial SSE line across chunk boundaries
    let mut feed = |bytes: &[u8], pending: &mut Vec<u8>, on_event: &mut F| {
        pending.extend_from_slice(bytes);
        while let Some(nl) = pending.iter().position(|&b| b == b'\n') {
            let line: Vec<u8> = pending.drain(..=nl).collect();
            let s = String::from_utf8_lossy(&line);
            if let Some(p) = s.trim_end().strip_prefix("data:") { on_event(p.trim()) }
        }
    };

    if chunked {
        loop {
            let mut sz = String::new();
            if r.read_line(&mut sz)? == 0 { break }
            let n = usize::from_str_radix(sz.trim().split(';').next().unwrap_or("0").trim(), 16)
                .unwrap_or(0);
            if n == 0 { let mut crlf = String::new(); let _ = r.read_line(&mut crlf); break }
            let mut buf = vec![0u8; n];
            r.read_exact(&mut buf)?;
            feed(&buf, &mut pending, &mut on_event);
            body.extend_from_slice(&buf);
            let mut crlf = [0u8; 2];
            r.read_exact(&mut crlf)?;   // trailing CRLF after each chunk
        }
    } else if let Some(n) = clen {
        body = vec![0u8; n];
        r.read_exact(&mut body)?;
        feed(&body.clone(), &mut pending, &mut on_event);
    } else {
        r.read_to_end(&mut body)?;
        feed(&body.clone(), &mut pending, &mut on_event);
    }
    Ok(body)
}

fn model_id(host: &str, port: u16) -> String {
    let s = TcpStream::connect((host, port)).expect("connect /v1/models");
    let mut w = s.try_clone().unwrap();
    write!(w, "GET /v1/models HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n").unwrap();
    let mut r = BufReader::new(s);
    let body = read_response(&mut r, |_| {}).unwrap();
    let v: serde_json::Value = serde_json::from_slice(&body).expect("models json");
    v["data"][0]["id"].as_str().unwrap_or("model").to_string()
}

struct Run { tokens: usize, ttft: Option<f64> }

fn one(host: &str, port: u16, model: &str, ids: &str, ngen: usize, stream: bool, t0: Instant) -> Run {
    let body = format!(
        "{{\"model\":\"{model}\",\"prompt\":[{ids}],\"max_tokens\":{ngen},\"temperature\":0,\
         \"ignore_eos\":true,\"stream\":{stream}}}"
    );
    let s = TcpStream::connect((host, port)).expect("connect");
    s.set_nodelay(true).ok();
    let mut w = s.try_clone().unwrap();
    write!(w, "POST /v1/completions HTTP/1.1\r\nHost: {host}\r\nContent-Type: application/json\r\n\
               Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).unwrap();
    w.flush().unwrap();

    let mut r = BufReader::new(s);
    let mut n = 0usize;
    let mut ttft = None;
    let raw = read_response(&mut r, |p| {
        if p == "[DONE]" { return }
        if ttft.is_none() { ttft = Some(t0.elapsed().as_secs_f64()) }
        n += 1;
    }).expect("response");

    if !stream {
        // No per-token events: take the server's own count, falling back to the request budget.
        let v: serde_json::Value = serde_json::from_slice(&raw).unwrap_or(serde_json::Value::Null);
        n = v["usage"]["completion_tokens"].as_u64().unwrap_or(ngen as u64) as usize;
        ttft = Some(t0.elapsed().as_secs_f64());
    }
    Run { tokens: n, ttft }
}

fn main() {
    let a: Vec<String> = std::env::args().collect();
    if a.len() < 5 {
        eprintln!("usage: loadgen <port> <ngen> <M> <ids-file> [--nostream] [--host H]");
        std::process::exit(2);
    }
    let port: u16 = a[1].parse().expect("port");
    let ngen: usize = a[2].parse().expect("ngen");
    let m: usize = a[3].parse().expect("M");
    let host = arg_val(&a, "--host").unwrap_or_else(|| "127.0.0.1".into());
    let stream = !arg_flag(&a, "--nostream");
    // PROMPT SET, not one prompt. A file with several non-empty lines is a SET: stream i gets
    // line i % lines.len(). Sending every slot the same prompt is the easy case for both
    // engines and actively misleading for MoE — identical tokens route to identical experts,
    // so ~8 of 128 activate instead of ~125 and the expert-weight traffic a real batch pays
    // never happens. Varied lengths also stop one uniform prefill from hiding scheduler
    // behaviour. A single-line file keeps the old behaviour exactly.
    let raw = std::fs::read_to_string(&a[4]).expect("ids file");
    let set: Vec<String> = raw.lines().map(|l| l.trim()).filter(|l| !l.is_empty())
        .map(|l| l.to_string()).collect();
    assert!(!set.is_empty(), "ids file is empty");
    if set.len() > 1 {
        let lens: Vec<usize> = set.iter().map(|p| p.split(',').count()).collect();
        eprintln!("prompt set: {} prompts, lengths {}..{}",
                  set.len(), lens.iter().min().unwrap(), lens.iter().max().unwrap());
    }

    let model = model_id(&host, port);
    // Warm exactly as the Python harness did: one full request, discarded.
    one(&host, port, &model, &set[0], ngen, stream, Instant::now());

    let t0 = Instant::now();
    let runs: Vec<Run> = std::thread::scope(|sc| {
        let hs: Vec<_> = (0..m).map(|i| {
            let (h2, mo, id2) = (host.clone(), model.clone(), set[i % set.len()].clone());
            sc.spawn(move || one(&h2, port, &mo, &id2, ngen, stream, t0))
        }).collect();
        hs.into_iter().map(|h| h.join().expect("stream panicked")).collect()
    });
    let dt = t0.elapsed().as_secs_f64();

    let total: usize = runs.iter().map(|r| r.tokens).sum();
    let mut tt: Vec<f64> = runs.iter().filter_map(|r| r.ttft).collect();
    tt.sort_by(|x, y| x.partial_cmp(y).unwrap());
    let q = |p: f64| -> f64 {
        if tt.is_empty() { f64::NAN } else { tt[((tt.len() - 1) as f64 * p) as usize] * 1e3 }
    };
    println!("RUST M={m}{}: {total} tokens in {dt:.3}s = AGGREGATE {:.1} tok/s ({:.1} per seq), \
              TTFT min {:.1} p50 {:.1} p90 {:.1} max {:.1} ms",
             if stream { "" } else { " nostream" }, total as f64 / dt, total as f64 / dt / m as f64,
             q(0.0), q(0.5), q(0.9), q(1.0));
}