613 lines
19 KiB
Rust
613 lines
19 KiB
Rust
//! 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);
|
|
}
|
|
}
|