Initial commit

crates/production/rtx-monitoring/src/metrics.rs

@@ -0,0 +1,612 @@
+//! Metrics collection and management.
+//!
+//! This module provides Prometheus-compatible metrics for monitoring
+//! RustyTorch++ ML inference and training workloads.
+
+use crate::{MonitoringError, MonitoringResult};
+use prometheus::{
+    Counter, CounterVec, Gauge, GaugeVec, Histogram, HistogramOpts, HistogramVec, Opts, Registry,
+};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+/// Metric type enumeration.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum MetricType {
+    /// Monotonically increasing counter
+    Counter,
+    /// Value that can go up and down
+    Gauge,
+    /// Distribution of values
+    Histogram,
+}
+
+/// Custom metric definition.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CustomMetric {
+    /// Metric name
+    pub name: String,
+    /// Help text
+    pub help: String,
+    /// Type of metric
+    pub metric_type: MetricType,
+    /// Label key-value pairs
+    pub labels: HashMap<String, String>,
+}
+
+/// Metric registry for managing Prometheus metrics.
+#[derive(Debug)]
+pub struct MetricRegistry {
+    registry: Registry,
+    counters: HashMap<String, Counter>,
+    gauges: HashMap<String, Gauge>,
+    histograms: HashMap<String, Histogram>,
+}
+
+impl MetricRegistry {
+    /// Create a new metric registry.
+    pub fn new() -> Self {
+        Self {
+            registry: Registry::new(),
+            counters: HashMap::new(),
+            gauges: HashMap::new(),
+            histograms: HashMap::new(),
+        }
+    }
+
+    /// Register a counter metric.
+    pub fn register_counter(&mut self, name: &str, help: &str) -> MonitoringResult<()> {
+        let counter = Counter::new(name, help)?;
+        self.registry.register(Box::new(counter.clone()))?;
+        self.counters.insert(name.to_string(), counter);
+        Ok(())
+    }
+
+    /// Register a gauge metric.
+    pub fn register_gauge(&mut self, name: &str, help: &str) -> MonitoringResult<()> {
+        let gauge = Gauge::new(name, help)?;
+        self.registry.register(Box::new(gauge.clone()))?;
+        self.gauges.insert(name.to_string(), gauge);
+        Ok(())
+    }
+
+    /// Register a histogram metric.
+    pub fn register_histogram(
+        &mut self,
+        name: &str,
+        help: &str,
+        buckets: Vec<f64>,
+    ) -> MonitoringResult<()> {
+        let opts = HistogramOpts::new(name, help).buckets(buckets);
+        let histogram = Histogram::with_opts(opts)?;
+        self.registry.register(Box::new(histogram.clone()))?;
+        self.histograms.insert(name.to_string(), histogram);
+        Ok(())
+    }
+
+    /// Increment a counter.
+    pub fn increment_counter(&self, name: &str) -> MonitoringResult<()> {
+        if let Some(counter) = self.counters.get(name) {
+            counter.inc();
+            Ok(())
+        } else {
+            Err(MonitoringError::metrics_error(format!(
+                "Counter '{name}' not found"
+            )))
+        }
+    }
+
+    /// Add to a counter.
+    pub fn add_counter(&self, name: &str, value: f64) -> MonitoringResult<()> {
+        if let Some(counter) = self.counters.get(name) {
+            counter.inc_by(value);
+            Ok(())
+        } else {
+            Err(MonitoringError::metrics_error(format!(
+                "Counter '{name}' not found"
+            )))
+        }
+    }
+
+    /// Set a gauge value.
+    pub fn set_gauge(&self, name: &str, value: f64) -> MonitoringResult<()> {
+        if let Some(gauge) = self.gauges.get(name) {
+            gauge.set(value);
+            Ok(())
+        } else {
+            Err(MonitoringError::metrics_error(format!(
+                "Gauge '{name}' not found"
+            )))
+        }
+    }
+
+    /// Observe a histogram value.
+    pub fn observe_histogram(&self, name: &str, value: f64) -> MonitoringResult<()> {
+        if let Some(histogram) = self.histograms.get(name) {
+            histogram.observe(value);
+            Ok(())
+        } else {
+            Err(MonitoringError::metrics_error(format!(
+                "Histogram '{name}' not found"
+            )))
+        }
+    }
+
+    /// Get the underlying registry.
+    pub fn registry(&self) -> &Registry {
+        &self.registry
+    }
+}
+
+impl Default for MetricRegistry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl From<prometheus::Error> for MonitoringError {
+    fn from(err: prometheus::Error) -> Self {
+        Self::metrics_error(err.to_string())
+    }
+}
+
+// ============================================================================
+// ML-Specific Metrics
+// ============================================================================
+
+/// Default latency buckets for inference (in seconds)
+pub const INFERENCE_LATENCY_BUCKETS: &[f64] = &[
+    0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
+];
+
+/// Default batch size buckets
+pub const BATCH_SIZE_BUCKETS: &[f64] = &[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0];
+
+/// ML inference metrics collection
+#[derive(Debug)]
+pub struct InferenceMetrics {
+    /// Total inference requests
+    pub requests_total: CounterVec,
+    /// Total inference errors
+    pub errors_total: CounterVec,
+    /// Inference latency histogram
+    pub latency_seconds: HistogramVec,
+    /// Tokens generated (for LLMs)
+    pub tokens_generated: CounterVec,
+    /// Tokens per second (throughput)
+    pub tokens_per_second: GaugeVec,
+    /// Batch size distribution
+    pub batch_size: HistogramVec,
+    /// Active requests gauge
+    pub active_requests: GaugeVec,
+    /// Queue depth
+    pub queue_depth: GaugeVec,
+    /// Cache hit rate
+    pub cache_hits: CounterVec,
+    /// Cache misses
+    pub cache_misses: CounterVec,
+    /// GPU memory usage
+    pub gpu_memory_bytes: GaugeVec,
+    /// GPU utilization percentage
+    pub gpu_utilization: GaugeVec,
+    /// Model load time
+    pub model_load_seconds: HistogramVec,
+    /// First token latency (time to first token for streaming)
+    pub time_to_first_token: HistogramVec,
+    /// Registry reference
+    registry: Registry,
+}
+
+impl InferenceMetrics {
+    /// Create a new inference metrics instance with default configuration
+    pub fn new() -> MonitoringResult<Self> {
+        let registry = Registry::new();
+
+        // Request counter with model and status labels
+        let requests_total = CounterVec::new(
+            Opts::new("rtx_inference_requests_total", "Total inference requests"),
+            &["model", "status"],
+        )?;
+        registry.register(Box::new(requests_total.clone()))?;
+
+        // Error counter with model and error_type labels
+        let errors_total = CounterVec::new(
+            Opts::new("rtx_inference_errors_total", "Total inference errors"),
+            &["model", "error_type"],
+        )?;
+        registry.register(Box::new(errors_total.clone()))?;
+
+        // Latency histogram
+        let latency_seconds = HistogramVec::new(
+            HistogramOpts::new(
+                "rtx_inference_latency_seconds",
+                "Inference latency in seconds",
+            )
+            .buckets(INFERENCE_LATENCY_BUCKETS.to_vec()),
+            &["model", "batch_size"],
+        )?;
+        registry.register(Box::new(latency_seconds.clone()))?;
+
+        // Tokens generated
+        let tokens_generated = CounterVec::new(
+            Opts::new("rtx_tokens_generated_total", "Total tokens generated"),
+            &["model"],
+        )?;
+        registry.register(Box::new(tokens_generated.clone()))?;
+
+        // Tokens per second gauge
+        let tokens_per_second = GaugeVec::new(
+            Opts::new("rtx_tokens_per_second", "Token generation throughput"),
+            &["model"],
+        )?;
+        registry.register(Box::new(tokens_per_second.clone()))?;
+
+        // Batch size histogram
+        let batch_size = HistogramVec::new(
+            HistogramOpts::new("rtx_batch_size", "Inference batch size distribution")
+                .buckets(BATCH_SIZE_BUCKETS.to_vec()),
+            &["model"],
+        )?;
+        registry.register(Box::new(batch_size.clone()))?;
+
+        // Active requests gauge
+        let active_requests = GaugeVec::new(
+            Opts::new("rtx_active_requests", "Currently active inference requests"),
+            &["model"],
+        )?;
+        registry.register(Box::new(active_requests.clone()))?;
+
+        // Queue depth gauge
+        let queue_depth = GaugeVec::new(
+            Opts::new("rtx_queue_depth", "Number of requests waiting in queue"),
+            &["model"],
+        )?;
+        registry.register(Box::new(queue_depth.clone()))?;
+
+        // Cache metrics
+        let cache_hits = CounterVec::new(
+            Opts::new("rtx_cache_hits_total", "KV cache hits"),
+            &["model", "cache_type"],
+        )?;
+        registry.register(Box::new(cache_hits.clone()))?;
+
+        let cache_misses = CounterVec::new(
+            Opts::new("rtx_cache_misses_total", "KV cache misses"),
+            &["model", "cache_type"],
+        )?;
+        registry.register(Box::new(cache_misses.clone()))?;
+
+        // GPU metrics
+        let gpu_memory_bytes = GaugeVec::new(
+            Opts::new("rtx_gpu_memory_bytes", "GPU memory usage in bytes"),
+            &["device", "type"],
+        )?;
+        registry.register(Box::new(gpu_memory_bytes.clone()))?;
+
+        let gpu_utilization = GaugeVec::new(
+            Opts::new("rtx_gpu_utilization_percent", "GPU utilization percentage"),
+            &["device"],
+        )?;
+        registry.register(Box::new(gpu_utilization.clone()))?;
+
+        // Model load time
+        let model_load_seconds = HistogramVec::new(
+            HistogramOpts::new("rtx_model_load_seconds", "Time to load model")
+                .buckets(vec![0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0]),
+            &["model"],
+        )?;
+        registry.register(Box::new(model_load_seconds.clone()))?;
+
+        // Time to first token
+        let time_to_first_token = HistogramVec::new(
+            HistogramOpts::new(
+                "rtx_time_to_first_token_seconds",
+                "Time to generate first token",
+            )
+            .buckets(vec![0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0]),
+            &["model"],
+        )?;
+        registry.register(Box::new(time_to_first_token.clone()))?;
+
+        Ok(Self {
+            requests_total,
+            errors_total,
+            latency_seconds,
+            tokens_generated,
+            tokens_per_second,
+            batch_size,
+            active_requests,
+            queue_depth,
+            cache_hits,
+            cache_misses,
+            gpu_memory_bytes,
+            gpu_utilization,
+            model_load_seconds,
+            time_to_first_token,
+            registry,
+        })
+    }
+
+    /// Record a successful inference request
+    pub fn record_request(&self, model: &str, latency_secs: f64, batch_size_val: usize) {
+        self.requests_total
+            .with_label_values(&[model, "success"])
+            .inc();
+        self.latency_seconds
+            .with_label_values(&[model, &batch_size_val.to_string()])
+            .observe(latency_secs);
+        self.batch_size
+            .with_label_values(&[model])
+            .observe(batch_size_val as f64);
+    }
+
+    /// Record a failed inference request
+    pub fn record_error(&self, model: &str, error_type: &str) {
+        self.requests_total
+            .with_label_values(&[model, "error"])
+            .inc();
+        self.errors_total
+            .with_label_values(&[model, error_type])
+            .inc();
+    }
+
+    /// Record tokens generated
+    pub fn record_tokens(&self, model: &str, count: u64) {
+        self.tokens_generated
+            .with_label_values(&[model])
+            .inc_by(count as f64);
+    }
+
+    /// Update tokens per second
+    pub fn update_throughput(&self, model: &str, tokens_per_sec: f64) {
+        self.tokens_per_second
+            .with_label_values(&[model])
+            .set(tokens_per_sec);
+    }
+
+    /// Update active request count
+    pub fn set_active_requests(&self, model: &str, count: i64) {
+        self.active_requests
+            .with_label_values(&[model])
+            .set(count as f64);
+    }
+
+    /// Update queue depth
+    pub fn set_queue_depth(&self, model: &str, depth: i64) {
+        self.queue_depth
+            .with_label_values(&[model])
+            .set(depth as f64);
+    }
+
+    /// Record cache hit
+    pub fn record_cache_hit(&self, model: &str, cache_type: &str) {
+        self.cache_hits
+            .with_label_values(&[model, cache_type])
+            .inc();
+    }
+
+    /// Record cache miss
+    pub fn record_cache_miss(&self, model: &str, cache_type: &str) {
+        self.cache_misses
+            .with_label_values(&[model, cache_type])
+            .inc();
+    }
+
+    /// Update GPU memory usage
+    pub fn set_gpu_memory(&self, device: &str, used_bytes: u64, total_bytes: u64) {
+        self.gpu_memory_bytes
+            .with_label_values(&[device, "used"])
+            .set(used_bytes as f64);
+        self.gpu_memory_bytes
+            .with_label_values(&[device, "total"])
+            .set(total_bytes as f64);
+    }
+
+    /// Update GPU utilization
+    pub fn set_gpu_utilization(&self, device: &str, utilization_percent: f64) {
+        self.gpu_utilization
+            .with_label_values(&[device])
+            .set(utilization_percent);
+    }
+
+    /// Record model load time
+    pub fn record_model_load(&self, model: &str, load_time_secs: f64) {
+        self.model_load_seconds
+            .with_label_values(&[model])
+            .observe(load_time_secs);
+    }
+
+    /// Record time to first token
+    pub fn record_ttft(&self, model: &str, ttft_secs: f64) {
+        self.time_to_first_token
+            .with_label_values(&[model])
+            .observe(ttft_secs);
+    }
+
+    /// Get the Prometheus registry
+    pub fn registry(&self) -> &Registry {
+        &self.registry
+    }
+
+    /// Export metrics in Prometheus format
+    pub fn export(&self) -> String {
+        use prometheus::Encoder;
+        let encoder = prometheus::TextEncoder::new();
+        let metric_families = self.registry.gather();
+        let mut buffer = Vec::new();
+        encoder
+            .encode(&metric_families, &mut buffer)
+            .unwrap_or_default();
+        String::from_utf8(buffer).unwrap_or_default()
+    }
+}
+
+impl Default for InferenceMetrics {
+    fn default() -> Self {
+        Self::new().expect("Failed to create default inference metrics")
+    }
+}
+
+/// Training metrics collection
+#[derive(Debug)]
+pub struct TrainingMetrics {
+    /// Training steps completed
+    pub steps_total: CounterVec,
+    /// Training loss
+    pub loss: GaugeVec,
+    /// Learning rate
+    pub learning_rate: GaugeVec,
+    /// Gradient norm
+    pub gradient_norm: GaugeVec,
+    /// Samples processed per second
+    pub samples_per_second: GaugeVec,
+    /// Epoch progress
+    pub epoch: GaugeVec,
+    /// Checkpoint save time
+    pub checkpoint_save_seconds: HistogramVec,
+    /// Registry
+    registry: Registry,
+}
+
+impl TrainingMetrics {
+    /// Create new training metrics
+    pub fn new() -> MonitoringResult<Self> {
+        let registry = Registry::new();
+
+        let steps_total = CounterVec::new(
+            Opts::new("rtx_training_steps_total", "Total training steps"),
+            &["model", "phase"],
+        )?;
+        registry.register(Box::new(steps_total.clone()))?;
+
+        let loss = GaugeVec::new(
+            Opts::new("rtx_training_loss", "Current training loss"),
+            &["model", "loss_type"],
+        )?;
+        registry.register(Box::new(loss.clone()))?;
+
+        let learning_rate = GaugeVec::new(
+            Opts::new("rtx_learning_rate", "Current learning rate"),
+            &["model"],
+        )?;
+        registry.register(Box::new(learning_rate.clone()))?;
+
+        let gradient_norm =
+            GaugeVec::new(Opts::new("rtx_gradient_norm", "Gradient norm"), &["model"])?;
+        registry.register(Box::new(gradient_norm.clone()))?;
+
+        let samples_per_second = GaugeVec::new(
+            Opts::new("rtx_samples_per_second", "Training throughput"),
+            &["model"],
+        )?;
+        registry.register(Box::new(samples_per_second.clone()))?;
+
+        let epoch = GaugeVec::new(Opts::new("rtx_training_epoch", "Current epoch"), &["model"])?;
+        registry.register(Box::new(epoch.clone()))?;
+
+        let checkpoint_save_seconds = HistogramVec::new(
+            HistogramOpts::new("rtx_checkpoint_save_seconds", "Checkpoint save time")
+                .buckets(vec![1.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0]),
+            &["model"],
+        )?;
+        registry.register(Box::new(checkpoint_save_seconds.clone()))?;
+
+        Ok(Self {
+            steps_total,
+            loss,
+            learning_rate,
+            gradient_norm,
+            samples_per_second,
+            epoch,
+            checkpoint_save_seconds,
+            registry,
+        })
+    }
+
+    /// Record a training step
+    pub fn record_step(&self, model: &str, phase: &str, loss_val: f64, lr: f64, grad_norm: f64) {
+        self.steps_total.with_label_values(&[model, phase]).inc();
+        self.loss.with_label_values(&[model, "total"]).set(loss_val);
+        self.learning_rate.with_label_values(&[model]).set(lr);
+        self.gradient_norm
+            .with_label_values(&[model])
+            .set(grad_norm);
+    }
+
+    /// Update throughput
+    pub fn update_throughput(&self, model: &str, samples_per_sec: f64) {
+        self.samples_per_second
+            .with_label_values(&[model])
+            .set(samples_per_sec);
+    }
+
+    /// Set current epoch
+    pub fn set_epoch(&self, model: &str, epoch_num: f64) {
+        self.epoch.with_label_values(&[model]).set(epoch_num);
+    }
+
+    /// Record checkpoint save time
+    pub fn record_checkpoint(&self, model: &str, save_time_secs: f64) {
+        self.checkpoint_save_seconds
+            .with_label_values(&[model])
+            .observe(save_time_secs);
+    }
+
+    /// Get registry
+    pub fn registry(&self) -> &Registry {
+        &self.registry
+    }
+}
+
+impl Default for TrainingMetrics {
+    fn default() -> Self {
+        Self::new().expect("Failed to create default training metrics")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_metric_registry() {
+        let mut registry = MetricRegistry::new();
+        assert!(
+            registry
+                .register_counter("test_counter", "Test counter")
+                .is_ok()
+        );
+        assert!(registry.register_gauge("test_gauge", "Test gauge").is_ok());
+        assert!(registry.increment_counter("test_counter").is_ok());
+        assert!(registry.set_gauge("test_gauge", 42.0).is_ok());
+    }
+
+    #[test]
+    fn test_inference_metrics() {
+        let metrics = InferenceMetrics::new().unwrap();
+
+        // Record some metrics
+        metrics.record_request("gpt-2", 0.1, 8);
+        metrics.record_tokens("gpt-2", 100);
+        metrics.update_throughput("gpt-2", 500.0);
+        metrics.set_active_requests("gpt-2", 5);
+        metrics.record_cache_hit("gpt-2", "kv");
+        metrics.set_gpu_memory("cuda:0", 4_000_000_000, 8_000_000_000);
+        metrics.set_gpu_utilization("cuda:0", 75.0);
+
+        // Export should contain our metrics
+        let output = metrics.export();
+        assert!(output.contains("rtx_inference_requests_total"));
+        assert!(output.contains("rtx_tokens_generated_total"));
+    }
+
+    #[test]
+    fn test_training_metrics() {
+        let metrics = TrainingMetrics::new().unwrap();
+
+        metrics.record_step("bert", "train", 0.5, 0.001, 1.5);
+        metrics.update_throughput("bert", 1000.0);
+        metrics.set_epoch("bert", 3.0);
+        metrics.record_checkpoint("bert", 10.5);
+    }
+}