furnace-server 0.13.0

//! LLM metrics, measured at the server edge from each request's events, so they mean the same
//! for every engine: time to first token, inter-token latency, queue time, token counts, and
//! how requests ended.

use std::fmt::Write as _;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

use super::{Event, FinishReason, LlmEngine};
use kuebiko::{Histogram, Percentiles, label};

/// The counters and histograms of one engine
#[derive(Default)]
pub(crate) struct EdgeMetrics {
    /// Received to the first `Text`
    ttft: Histogram,
    /// Between consecutive `Text`s of one stream
    itl: Histogram,
    /// Received to `Started`
    queue: Histogram,
    prompt_tokens: AtomicU64,
    generation_tokens: AtomicU64,
    /// Tokens after the first, and the time they took, over finished streams
    decode_tokens: AtomicU64,
    decode_nanos: AtomicU64,
    waiting: AtomicU64,
    running: AtomicU64,
    finished_stop: AtomicU64,
    finished_length: AtomicU64,
    finished_cancelled: AtomicU64,
    finished_error: AtomicU64,
    /// Refused with 503 by the concurrency limit or the engine
    rejected: AtomicU64,
}

impl EdgeMetrics {
    pub(crate) fn reject(&self) {
        self.rejected.fetch_add(1, Ordering::Relaxed);
    }
}

/// How requests ended, per `finish_reason` (`cancelled`: the client left first)
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
pub struct FinishedCounts {
    pub stop: u64,
    pub length: u64,
    pub cancelled: u64,
    pub error: u64,
}

/// `GET /metrics` for one LLM
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LlmMetrics {
    pub name: String,
    pub requests_waiting: u64,
    pub requests_running: u64,
    pub requests_finished: FinishedCounts,
    /// Requests refused with 503
    pub requests_rejected: u64,
    pub prompt_tokens: u64,
    pub generation_tokens: u64,
    /// Decoding speed of one stream, over the streams that finished: tokens after the first,
    /// per second after the first
    pub tokens_per_second: Option<f64>,
    pub time_to_first_token_ms: Percentiles,
    pub inter_token_latency_ms: Percentiles,
    pub queue_time_ms: Percentiles,
    /// Furnace's scheduler only: model time for one prefill, and for one decode step
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub prefill_ms: Option<Percentiles>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decode_step_ms: Option<Percentiles>,
}

pub(crate) fn json(name: &str, metrics: &EdgeMetrics, engine: &dyn LlmEngine) -> LlmMetrics {
    let load = |c: &AtomicU64| c.load(Ordering::Relaxed);
    let decode_tokens = load(&metrics.decode_tokens);
    let decode_seconds = load(&metrics.decode_nanos) as f64 / 1e9;
    let stats = engine.scheduler_stats();
    LlmMetrics {
        name: name.to_string(),
        requests_waiting: load(&metrics.waiting),
        requests_running: load(&metrics.running),
        requests_finished: FinishedCounts {
            stop: load(&metrics.finished_stop),
            length: load(&metrics.finished_length),
            cancelled: load(&metrics.finished_cancelled),
            error: load(&metrics.finished_error),
        },
        requests_rejected: load(&metrics.rejected),
        prompt_tokens: load(&metrics.prompt_tokens),
        generation_tokens: load(&metrics.generation_tokens),
        tokens_per_second: (decode_seconds > 0.0).then(|| decode_tokens as f64 / decode_seconds),
        time_to_first_token_ms: metrics.ttft.percentiles_ms(),
        inter_token_latency_ms: metrics.itl.percentiles_ms(),
        queue_time_ms: metrics.queue.percentiles_ms(),
        prefill_ms: stats.map(|s| s.prefill.percentiles_ms()),
        decode_step_ms: stats.map(|s| s.decode_step.percentiles_ms()),
    }
}

/// Append every LLM's series, labelled `model`, in the Prometheus text format
pub(crate) fn render(out: &mut String, engines: &[(&str, &EdgeMetrics, &dyn LlmEngine)]) {
    if engines.is_empty() {
        return;
    }
    type Value = fn(&EdgeMetrics) -> &AtomicU64;
    let values: [(&str, &str, &str, Value); 5] = [
        (
            "furnace_llm_prompt_tokens_total",
            "counter",
            "Prompt tokens of the requests the engine started",
            |m| &m.prompt_tokens,
        ),
        (
            "furnace_llm_generation_tokens_total",
            "counter",
            "Tokens generated by finished requests; tokens/s is its rate()",
            |m| &m.generation_tokens,
        ),
        (
            "furnace_llm_requests_running",
            "gauge",
            "Requests the engine is generating for",
            |m| &m.running,
        ),
        (
            "furnace_llm_requests_waiting",
            "gauge",
            "Requests waiting for the engine to start them",
            |m| &m.waiting,
        ),
        (
            "furnace_llm_requests_rejected_total",
            "counter",
            "Requests refused with 503 at the concurrency limit or by a full engine",
            |m| &m.rejected,
        ),
    ];
    for (name, kind, help, value) in values {
        let _ = writeln!(out, "# HELP {name} {help}");
        let _ = writeln!(out, "# TYPE {name} {kind}");
        for (model, metrics, _) in engines {
            let _ = writeln!(
                out,
                "{name}{{model=\"{}\"}} {}",
                label(model),
                value(metrics).load(Ordering::Relaxed)
            );
        }
    }
    let name = "furnace_llm_requests_finished_total";
    let _ = writeln!(out, "# HELP {name} Requests by how they ended");
    let _ = writeln!(out, "# TYPE {name} counter");
    for (model, m, _) in engines {
        for (reason, count) in [
            ("stop", &m.finished_stop),
            ("length", &m.finished_length),
            ("cancelled", &m.finished_cancelled),
            ("error", &m.finished_error),
        ] {
            let _ = writeln!(
                out,
                "{name}{{model=\"{}\",reason=\"{reason}\"}} {}",
                label(model),
                count.load(Ordering::Relaxed)
            );
        }
    }
    type Series = for<'x> fn(&'x EdgeMetrics, &'x dyn LlmEngine) -> Option<&'x Histogram>;
    fn ttft<'x>(m: &'x EdgeMetrics, _: &'x dyn LlmEngine) -> Option<&'x Histogram> {
        Some(&m.ttft)
    }
    fn itl<'x>(m: &'x EdgeMetrics, _: &'x dyn LlmEngine) -> Option<&'x Histogram> {
        Some(&m.itl)
    }
    fn queue<'x>(m: &'x EdgeMetrics, _: &'x dyn LlmEngine) -> Option<&'x Histogram> {
        Some(&m.queue)
    }
    fn prefill<'x>(_: &'x EdgeMetrics, e: &'x dyn LlmEngine) -> Option<&'x Histogram> {
        e.scheduler_stats().map(|s| &s.prefill)
    }
    fn decode_step<'x>(_: &'x EdgeMetrics, e: &'x dyn LlmEngine) -> Option<&'x Histogram> {
        e.scheduler_stats().map(|s| &s.decode_step)
    }
    let histograms: [(&str, &str, Series); 5] = [
        (
            "furnace_llm_time_to_first_token_seconds",
            "Request received to the first token sent",
            ttft,
        ),
        (
            "furnace_llm_inter_token_latency_seconds",
            "Between consecutive tokens of one stream",
            itl,
        ),
        (
            "furnace_llm_queue_seconds",
            "Request received to the engine starting on it",
            queue,
        ),
        (
            "furnace_llm_prefill_seconds",
            "Model time for one prompt's prefill (furnace's scheduler)",
            prefill,
        ),
        (
            "furnace_llm_decode_step_seconds",
            "Model time for one decode step over a worker's sequences (furnace's scheduler)",
            decode_step,
        ),
    ];
    for (name, help, series) in histograms {
        let present: Vec<_> = engines
            .iter()
            .filter_map(|&(model, m, e)| series(m, e).map(|h| (model, h)))
            .collect();
        if present.is_empty() {
            continue;
        }
        let _ = writeln!(out, "# HELP {name} {help}");
        let _ = writeln!(out, "# TYPE {name} histogram");
        for (model, histogram) in present {
            histogram.render_series(out, name, &format!("model=\"{}\"", label(model)));
        }
    }
}

/// A request's events, timed on the way out. Dropping it before the end counts as a
/// cancellation, and drops the receiver, which cancels the request in the engine.
pub(crate) struct Tracked {
    events: mpsc::Receiver<Event>,
    metrics: Arc<EdgeMetrics>,
    received: Instant,
    /// When the last `Text` went out
    last_text: Option<Instant>,
    first_text: Option<Instant>,
    state: State,
    pub(crate) prompt_tokens: usize,
}

#[derive(Clone, Copy, PartialEq)]
enum State {
    Waiting,
    Running,
    Ended,
}

impl Tracked {
    pub(crate) fn new(
        events: mpsc::Receiver<Event>,
        metrics: Arc<EdgeMetrics>,
        received: Instant,
    ) -> Self {
        metrics.waiting.fetch_add(1, Ordering::Relaxed);
        Self {
            events,
            metrics,
            received,
            last_text: None,
            first_text: None,
            state: State::Waiting,
            prompt_tokens: 0,
        }
    }

    /// The next event; `None` when the engine dropped the request without finishing it
    pub(crate) async fn next(&mut self) -> Option<Event> {
        let event = self.events.recv().await;
        let now = Instant::now();
        let m = &*self.metrics;
        match &event {
            Some(Event::Started { prompt_tokens }) => {
                self.prompt_tokens = *prompt_tokens;
                m.prompt_tokens
                    .fetch_add(*prompt_tokens as u64, Ordering::Relaxed);
                m.queue.observe(now - self.received);
                self.set_state(State::Running);
            }
            Some(Event::Text(_)) => {
                match self.last_text {
                    Some(last) => m.itl.observe(now - last),
                    None => {
                        m.ttft.observe(now - self.received);
                        self.first_text = Some(now);
                    }
                }
                self.last_text = Some(now);
            }
            Some(Event::Finished {
                reason,
                completion_tokens,
            }) => {
                m.generation_tokens
                    .fetch_add(*completion_tokens as u64, Ordering::Relaxed);
                if let (Some(first), Some(last)) = (self.first_text, self.last_text)
                    && *completion_tokens > 1
                    && last > first
                {
                    m.decode_tokens
                        .fetch_add(*completion_tokens as u64 - 1, Ordering::Relaxed);
                    m.decode_nanos.fetch_add(
                        u64::try_from((last - first).as_nanos()).unwrap_or(u64::MAX),
                        Ordering::Relaxed,
                    );
                }
                match reason {
                    FinishReason::Stop => &m.finished_stop,
                    FinishReason::Length => &m.finished_length,
                }
                .fetch_add(1, Ordering::Relaxed);
                self.set_state(State::Ended);
            }
            Some(Event::Error(_)) | None => {
                m.finished_error.fetch_add(1, Ordering::Relaxed);
                self.set_state(State::Ended);
            }
        }
        event
    }

    fn set_state(&mut self, state: State) {
        let m = &*self.metrics;
        match self.state {
            State::Waiting => m.waiting.fetch_sub(1, Ordering::Relaxed),
            State::Running => m.running.fetch_sub(1, Ordering::Relaxed),
            State::Ended => return,
        };
        if state == State::Running {
            m.running.fetch_add(1, Ordering::Relaxed);
        }
        self.state = state;
    }
}

impl Drop for Tracked {
    fn drop(&mut self) {
        if self.state != State::Ended {
            self.metrics
                .finished_cancelled
                .fetch_add(1, Ordering::Relaxed);
            self.set_state(State::Ended);
        }
    }
}