rustytorch/demos/inference-profiler-shared/src/metrics.rs

//! Performance metrics for inference profiling.

use serde::{Deserialize, Serialize};

/// Latency statistics from benchmark measurements.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct LatencyMetrics {
    /// Mean latency in milliseconds
    pub mean_ms: f64,
    /// Standard deviation in milliseconds
    pub std_ms: f64,
    /// Minimum latency in milliseconds
    pub min_ms: f64,
    /// Maximum latency in milliseconds
    pub max_ms: f64,
    /// 50th percentile (median) in milliseconds
    pub p50_ms: f64,
    /// 95th percentile in milliseconds
    pub p95_ms: f64,
    /// 99th percentile in milliseconds
    pub p99_ms: f64,
}

impl LatencyMetrics {
    /// Creates latency metrics from a collection of measurements.
    ///
    /// # Arguments
    /// * `measurements_ms` - Latency measurements in milliseconds
    ///
    /// # Errors
    /// Returns an error if measurements are empty or contain invalid values.
    pub fn from_measurements(measurements_ms: &[f64]) -> Result<Self, String> {
        if measurements_ms.is_empty() {
            return Err("measurements cannot be empty".to_string());
        }

        for &m in measurements_ms {
            if !m.is_finite() || m < 0.0 {
                return Err(format!("invalid measurement: {m}"));
            }
        }

        let mut sorted = measurements_ms.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());

        let mean = sorted.iter().sum::<f64>() / sorted.len() as f64;
        let variance =
            sorted.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / sorted.len() as f64;
        let std = variance.sqrt();

        let min = sorted[0];
        let max = sorted[sorted.len() - 1];

        let p50 = percentile(&sorted, 50.0);
        let p95 = percentile(&sorted, 95.0);
        let p99 = percentile(&sorted, 99.0);

        Ok(Self {
            mean_ms: mean,
            std_ms: std,
            min_ms: min,
            max_ms: max,
            p50_ms: p50,
            p95_ms: p95,
            p99_ms: p99,
        })
    }
}

/// Memory usage statistics.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MemoryMetrics {
    /// Peak memory usage in megabytes
    pub peak_memory_mb: f64,
    /// Currently allocated memory in megabytes
    pub allocated_mb: f64,
    /// Reserved memory in megabytes (for GPU)
    pub reserved_mb: f64,
}

impl MemoryMetrics {
    /// Creates new memory metrics.
    ///
    /// # Errors
    /// Returns an error if any value is negative or non-finite.
    pub fn new(peak_memory_mb: f64, allocated_mb: f64, reserved_mb: f64) -> Result<Self, String> {
        if !peak_memory_mb.is_finite() || peak_memory_mb < 0.0 {
            return Err(format!("invalid peak_memory_mb: {peak_memory_mb}"));
        }
        if !allocated_mb.is_finite() || allocated_mb < 0.0 {
            return Err(format!("invalid allocated_mb: {allocated_mb}"));
        }
        if !reserved_mb.is_finite() || reserved_mb < 0.0 {
            return Err(format!("invalid reserved_mb: {reserved_mb}"));
        }

        Ok(Self {
            peak_memory_mb,
            allocated_mb,
            reserved_mb,
        })
    }
}

/// Throughput statistics.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ThroughputMetrics {
    /// Samples processed per second
    pub samples_per_sec: f64,
    /// Tokens processed per second (optional, for sequence models)
    pub tokens_per_sec: Option<f64>,
}

impl ThroughputMetrics {
    /// Creates new throughput metrics.
    ///
    /// # Errors
    /// Returns an error if `samples_per_sec` is invalid.
    pub fn new(samples_per_sec: f64, tokens_per_sec: Option<f64>) -> Result<Self, String> {
        if !samples_per_sec.is_finite() || samples_per_sec < 0.0 {
            return Err(format!("invalid samples_per_sec: {samples_per_sec}"));
        }

        if let Some(tps) = tokens_per_sec
            && (!tps.is_finite() || tps < 0.0)
        {
            return Err(format!("invalid tokens_per_sec: {tps}"));
        }

        Ok(Self {
            samples_per_sec,
            tokens_per_sec,
        })
    }

    /// Calculates throughput from batch size and mean latency.
    ///
    /// # Arguments
    /// * `batch_size` - Number of samples in batch
    /// * `mean_latency_ms` - Mean latency in milliseconds
    ///
    /// # Errors
    /// Returns an error if inputs are invalid.
    pub fn from_latency(batch_size: usize, mean_latency_ms: f64) -> Result<Self, String> {
        if batch_size == 0 {
            return Err("batch_size cannot be zero".to_string());
        }
        if !mean_latency_ms.is_finite() || mean_latency_ms <= 0.0 {
            return Err(format!("invalid mean_latency_ms: {mean_latency_ms}"));
        }

        let samples_per_sec = (batch_size as f64 * 1000.0) / mean_latency_ms;
        Self::new(samples_per_sec, None)
    }
}

/// Computes the percentile of a sorted array.
fn percentile(sorted_data: &[f64], p: f64) -> f64 {
    if sorted_data.is_empty() {
        return 0.0;
    }

    if sorted_data.len() == 1 {
        return sorted_data[0];
    }

    let rank = (p / 100.0) * (sorted_data.len() - 1) as f64;
    let lower_idx = rank.floor() as usize;
    let upper_idx = rank.ceil() as usize;
    let weight = rank - lower_idx as f64;

    sorted_data[lower_idx] * (1.0 - weight) + sorted_data[upper_idx] * weight
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_latency_metrics_from_measurements_success() {
        let measurements = vec![10.0, 20.0, 15.0, 25.0, 30.0];
        let metrics = LatencyMetrics::from_measurements(&measurements);

        assert!(metrics.is_ok());
        let m = metrics.unwrap();
        assert_relative_eq!(m.mean_ms, 20.0, epsilon = 0.01);
        assert_relative_eq!(m.min_ms, 10.0, epsilon = 0.01);
        assert_relative_eq!(m.max_ms, 30.0, epsilon = 0.01);
    }

    #[test]
    fn test_latency_metrics_empty_measurements() {
        let measurements: Vec<f64> = vec![];
        let result = LatencyMetrics::from_measurements(&measurements);

        assert!(result.is_err());
        assert_eq!(result.unwrap_err(), "measurements cannot be empty");
    }

    #[test]
    fn test_latency_metrics_invalid_measurement_negative() {
        let measurements = vec![10.0, -5.0, 15.0];
        let result = LatencyMetrics::from_measurements(&measurements);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid measurement"));
    }

    #[test]
    fn test_latency_metrics_invalid_measurement_nan() {
        let measurements = vec![10.0, f64::NAN, 15.0];
        let result = LatencyMetrics::from_measurements(&measurements);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid measurement"));
    }

    #[test]
    fn test_latency_metrics_invalid_measurement_infinity() {
        let measurements = vec![10.0, f64::INFINITY, 15.0];
        let result = LatencyMetrics::from_measurements(&measurements);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid measurement"));
    }

    #[test]
    fn test_latency_metrics_percentile_calculations() {
        let measurements: Vec<f64> = (1..=100).map(|x| x as f64).collect();
        let metrics = LatencyMetrics::from_measurements(&measurements).unwrap();

        assert_relative_eq!(metrics.p50_ms, 50.0, epsilon = 1.0);
        assert_relative_eq!(metrics.p95_ms, 95.0, epsilon = 1.0);
        assert_relative_eq!(metrics.p99_ms, 99.0, epsilon = 1.0);
    }

    #[test]
    fn test_latency_metrics_single_measurement() {
        let measurements = vec![42.5];
        let metrics = LatencyMetrics::from_measurements(&measurements).unwrap();

        assert_relative_eq!(metrics.mean_ms, 42.5, epsilon = 0.01);
        assert_relative_eq!(metrics.std_ms, 0.0, epsilon = 0.01);
        assert_relative_eq!(metrics.min_ms, 42.5, epsilon = 0.01);
        assert_relative_eq!(metrics.max_ms, 42.5, epsilon = 0.01);
        assert_relative_eq!(metrics.p50_ms, 42.5, epsilon = 0.01);
        assert_relative_eq!(metrics.p95_ms, 42.5, epsilon = 0.01);
        assert_relative_eq!(metrics.p99_ms, 42.5, epsilon = 0.01);
    }

    #[test]
    fn test_latency_metrics_std_calculation() {
        let measurements = vec![2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
        let metrics = LatencyMetrics::from_measurements(&measurements).unwrap();

        assert_relative_eq!(metrics.mean_ms, 5.0, epsilon = 0.01);
        assert_relative_eq!(metrics.std_ms, 2.0, epsilon = 0.01);
    }

    #[test]
    fn test_latency_metrics_serialization() {
        let measurements = vec![10.0, 20.0, 30.0];
        let metrics = LatencyMetrics::from_measurements(&measurements).unwrap();
        let json = serde_json::to_string(&metrics).expect("serialization failed");

        assert!(json.contains("\"mean_ms\""));
        assert!(json.contains("\"std_ms\""));
    }

    #[test]
    fn test_memory_metrics_new_success() {
        let metrics = MemoryMetrics::new(100.0, 80.0, 120.0);

        assert!(metrics.is_ok());
        let m = metrics.unwrap();
        assert_relative_eq!(m.peak_memory_mb, 100.0, epsilon = 0.01);
        assert_relative_eq!(m.allocated_mb, 80.0, epsilon = 0.01);
        assert_relative_eq!(m.reserved_mb, 120.0, epsilon = 0.01);
    }

    #[test]
    fn test_memory_metrics_negative_peak() {
        let result = MemoryMetrics::new(-100.0, 80.0, 120.0);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid peak_memory_mb"));
    }

    #[test]
    fn test_memory_metrics_negative_allocated() {
        let result = MemoryMetrics::new(100.0, -80.0, 120.0);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid allocated_mb"));
    }

    #[test]
    fn test_memory_metrics_negative_reserved() {
        let result = MemoryMetrics::new(100.0, 80.0, -120.0);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid reserved_mb"));
    }

    #[test]
    fn test_memory_metrics_nan_values() {
        assert!(MemoryMetrics::new(f64::NAN, 80.0, 120.0).is_err());
        assert!(MemoryMetrics::new(100.0, f64::NAN, 120.0).is_err());
        assert!(MemoryMetrics::new(100.0, 80.0, f64::NAN).is_err());
    }

    #[test]
    fn test_memory_metrics_infinity_values() {
        assert!(MemoryMetrics::new(f64::INFINITY, 80.0, 120.0).is_err());
        assert!(MemoryMetrics::new(100.0, f64::INFINITY, 120.0).is_err());
        assert!(MemoryMetrics::new(100.0, 80.0, f64::INFINITY).is_err());
    }

    #[test]
    fn test_memory_metrics_zero_values() {
        let metrics = MemoryMetrics::new(0.0, 0.0, 0.0);
        assert!(metrics.is_ok());
    }

    #[test]
    fn test_memory_metrics_serialization() {
        let metrics = MemoryMetrics::new(100.0, 80.0, 120.0).unwrap();
        let json = serde_json::to_string(&metrics).expect("serialization failed");

        assert!(json.contains("\"peak_memory_mb\""));
        assert!(json.contains("\"allocated_mb\""));
        assert!(json.contains("\"reserved_mb\""));
    }

    #[test]
    fn test_throughput_metrics_new_success() {
        let metrics = ThroughputMetrics::new(100.0, Some(500.0));

        assert!(metrics.is_ok());
        let m = metrics.unwrap();
        assert_relative_eq!(m.samples_per_sec, 100.0, epsilon = 0.01);
        assert!(m.tokens_per_sec.is_some());
        assert_relative_eq!(m.tokens_per_sec.unwrap(), 500.0, epsilon = 0.01);
    }

    #[test]
    fn test_throughput_metrics_new_without_tokens() {
        let metrics = ThroughputMetrics::new(100.0, None);

        assert!(metrics.is_ok());
        let m = metrics.unwrap();
        assert_relative_eq!(m.samples_per_sec, 100.0, epsilon = 0.01);
        assert!(m.tokens_per_sec.is_none());
    }

    #[test]
    fn test_throughput_metrics_negative_samples() {
        let result = ThroughputMetrics::new(-100.0, None);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid samples_per_sec"));
    }

    #[test]
    fn test_throughput_metrics_negative_tokens() {
        let result = ThroughputMetrics::new(100.0, Some(-500.0));

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid tokens_per_sec"));
    }

    #[test]
    fn test_throughput_metrics_nan_samples() {
        let result = ThroughputMetrics::new(f64::NAN, None);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid samples_per_sec"));
    }

    #[test]
    fn test_throughput_metrics_infinity_tokens() {
        let result = ThroughputMetrics::new(100.0, Some(f64::INFINITY));

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid tokens_per_sec"));
    }

    #[test]
    fn test_throughput_from_latency_success() {
        let metrics = ThroughputMetrics::from_latency(8, 100.0);

        assert!(metrics.is_ok());
        let m = metrics.unwrap();
        assert_relative_eq!(m.samples_per_sec, 80.0, epsilon = 0.01);
    }

    #[test]
    fn test_throughput_from_latency_zero_batch_size() {
        let result = ThroughputMetrics::from_latency(0, 100.0);

        assert!(result.is_err());
        assert_eq!(result.unwrap_err(), "batch_size cannot be zero");
    }

    #[test]
    fn test_throughput_from_latency_zero_latency() {
        let result = ThroughputMetrics::from_latency(8, 0.0);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid mean_latency_ms"));
    }

    #[test]
    fn test_throughput_from_latency_negative_latency() {
        let result = ThroughputMetrics::from_latency(8, -100.0);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("invalid mean_latency_ms"));
    }

    #[test]
    fn test_throughput_from_latency_calculation() {
        let metrics = ThroughputMetrics::from_latency(32, 50.0).unwrap();
        assert_relative_eq!(metrics.samples_per_sec, 640.0, epsilon = 0.1);
    }

    #[test]
    fn test_throughput_metrics_serialization() {
        let metrics = ThroughputMetrics::new(100.0, Some(500.0)).unwrap();
        let json = serde_json::to_string(&metrics).expect("serialization failed");

        assert!(json.contains("\"samples_per_sec\""));
        assert!(json.contains("\"tokens_per_sec\""));
    }

    #[test]
    fn test_percentile_empty_array() {
        let data: Vec<f64> = vec![];
        let p50 = percentile(&data, 50.0);
        assert_relative_eq!(p50, 0.0, epsilon = 0.01);
    }

    #[test]
    fn test_percentile_single_element() {
        let data = vec![42.0];
        let p50 = percentile(&data, 50.0);
        assert_relative_eq!(p50, 42.0, epsilon = 0.01);
    }

    #[test]
    fn test_percentile_interpolation() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let p50 = percentile(&data, 50.0);
        assert_relative_eq!(p50, 3.0, epsilon = 0.01);

        let p75 = percentile(&data, 75.0);
        assert_relative_eq!(p75, 4.0, epsilon = 0.01);
    }

    #[test]
    fn test_percentile_edge_cases() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];

        let p0 = percentile(&data, 0.0);
        assert_relative_eq!(p0, 1.0, epsilon = 0.01);

        let p100 = percentile(&data, 100.0);
        assert_relative_eq!(p100, 5.0, epsilon = 0.01);
    }
}