Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,640 @@
|
||||
//! Alert management and notification system.
|
||||
//!
|
||||
//! This module provides alerting capabilities for RustyTorch++ including
|
||||
//! alert rule evaluation, notification dispatch, and alert history tracking.
|
||||
|
||||
use crate::MonitoringResult;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use parking_lot::RwLock;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Alert severity levels.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub enum AlertSeverity {
|
||||
/// Critical alert - immediate action required
|
||||
Critical,
|
||||
/// Warning - attention needed soon
|
||||
Warning,
|
||||
/// Info - informational alert
|
||||
Info,
|
||||
}
|
||||
|
||||
impl AlertSeverity {
|
||||
/// Get string representation
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Critical => "critical",
|
||||
Self::Warning => "warning",
|
||||
Self::Info => "info",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Alert state for tracking firing status
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum AlertState {
|
||||
/// Alert is currently firing
|
||||
Firing,
|
||||
/// Alert was firing but is now resolved
|
||||
Resolved,
|
||||
/// Alert is pending (condition met but waiting for duration)
|
||||
Pending,
|
||||
}
|
||||
|
||||
/// Alert rule definition.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AlertRule {
|
||||
/// Unique rule name
|
||||
pub name: String,
|
||||
/// Alert severity
|
||||
pub severity: AlertSeverity,
|
||||
/// Human-readable description
|
||||
pub description: String,
|
||||
/// Metric name to evaluate
|
||||
pub metric: String,
|
||||
/// Threshold condition
|
||||
pub condition: AlertCondition,
|
||||
/// Duration the condition must hold before firing
|
||||
pub duration: std::time::Duration,
|
||||
/// Labels to add to the alert
|
||||
pub labels: HashMap<String, String>,
|
||||
/// Annotations (additional metadata)
|
||||
pub annotations: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl AlertRule {
|
||||
/// Create a new alert rule
|
||||
pub fn new(
|
||||
name: &str,
|
||||
metric: &str,
|
||||
condition: AlertCondition,
|
||||
severity: AlertSeverity,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
severity,
|
||||
description: String::new(),
|
||||
metric: metric.to_string(),
|
||||
condition,
|
||||
duration: std::time::Duration::from_secs(0),
|
||||
labels: HashMap::new(),
|
||||
annotations: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set description
|
||||
pub fn with_description(mut self, desc: &str) -> Self {
|
||||
self.description = desc.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set duration before firing
|
||||
pub fn with_duration(mut self, duration: std::time::Duration) -> Self {
|
||||
self.duration = duration;
|
||||
self
|
||||
}
|
||||
|
||||
/// Add label
|
||||
pub fn with_label(mut self, key: &str, value: &str) -> Self {
|
||||
self.labels.insert(key.to_string(), value.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Add annotation
|
||||
pub fn with_annotation(mut self, key: &str, value: &str) -> Self {
|
||||
self.annotations.insert(key.to_string(), value.to_string());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Alert condition types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum AlertCondition {
|
||||
/// Value is greater than threshold
|
||||
GreaterThan(f64),
|
||||
/// Value is less than threshold
|
||||
LessThan(f64),
|
||||
/// Value is greater than or equal to threshold
|
||||
GreaterThanOrEqual(f64),
|
||||
/// Value is less than or equal to threshold
|
||||
LessThanOrEqual(f64),
|
||||
/// Value equals threshold
|
||||
Equal(f64),
|
||||
/// Value is not equal to threshold
|
||||
NotEqual(f64),
|
||||
/// Value is absent (no data)
|
||||
Absent,
|
||||
}
|
||||
|
||||
impl AlertCondition {
|
||||
/// Evaluate the condition against a value
|
||||
pub fn evaluate(&self, value: Option<f64>) -> bool {
|
||||
match (self, value) {
|
||||
(Self::GreaterThan(threshold), Some(v)) => v > *threshold,
|
||||
(Self::LessThan(threshold), Some(v)) => v < *threshold,
|
||||
(Self::GreaterThanOrEqual(threshold), Some(v)) => v >= *threshold,
|
||||
(Self::LessThanOrEqual(threshold), Some(v)) => v <= *threshold,
|
||||
(Self::Equal(threshold), Some(v)) => (v - threshold).abs() < f64::EPSILON,
|
||||
(Self::NotEqual(threshold), Some(v)) => (v - threshold).abs() >= f64::EPSILON,
|
||||
(Self::Absent, None) => true,
|
||||
(Self::Absent, Some(_)) => false,
|
||||
(_, None) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Notification configuration.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Notification {
|
||||
/// Unique notification ID
|
||||
pub id: String,
|
||||
/// Alert message
|
||||
pub message: String,
|
||||
/// Alert severity
|
||||
pub severity: AlertSeverity,
|
||||
/// When the notification was created
|
||||
pub timestamp: DateTime<Utc>,
|
||||
/// Alert state
|
||||
pub state: AlertState,
|
||||
/// Rule name that triggered this notification
|
||||
pub rule_name: String,
|
||||
/// Labels from the alert
|
||||
pub labels: HashMap<String, String>,
|
||||
/// Current metric value
|
||||
pub value: Option<f64>,
|
||||
}
|
||||
|
||||
impl Notification {
|
||||
/// Create a new notification from a rule
|
||||
pub fn from_rule(rule: &AlertRule, state: AlertState, value: Option<f64>) -> Self {
|
||||
Self {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
message: rule.description.clone(),
|
||||
severity: rule.severity,
|
||||
timestamp: Utc::now(),
|
||||
state,
|
||||
rule_name: rule.name.clone(),
|
||||
labels: rule.labels.clone(),
|
||||
value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Notification channel for dispatching alerts
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct NotificationChannel {
|
||||
/// Channel name
|
||||
pub name: String,
|
||||
/// Channel type
|
||||
pub channel_type: ChannelType,
|
||||
/// Channel configuration
|
||||
pub config: HashMap<String, String>,
|
||||
/// Minimum severity to send
|
||||
pub min_severity: AlertSeverity,
|
||||
}
|
||||
|
||||
/// Types of notification channels
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum ChannelType {
|
||||
/// Webhook (HTTP POST)
|
||||
Webhook,
|
||||
/// Email
|
||||
Email,
|
||||
/// Slack
|
||||
Slack,
|
||||
/// PagerDuty
|
||||
PagerDuty,
|
||||
/// Console/Log output
|
||||
Console,
|
||||
}
|
||||
|
||||
/// Active alert tracking
|
||||
#[derive(Debug, Clone)]
|
||||
struct ActiveAlert {
|
||||
rule: AlertRule,
|
||||
first_firing: DateTime<Utc>,
|
||||
last_evaluation: DateTime<Utc>,
|
||||
state: AlertState,
|
||||
value: Option<f64>,
|
||||
}
|
||||
|
||||
/// Alert manager for handling alerts and notifications.
|
||||
#[derive(Debug)]
|
||||
pub struct AlertManager {
|
||||
/// Registered alert rules
|
||||
rules: Arc<RwLock<HashMap<String, AlertRule>>>,
|
||||
/// Active alerts
|
||||
active_alerts: Arc<RwLock<HashMap<String, ActiveAlert>>>,
|
||||
/// Notification history
|
||||
notifications: Arc<RwLock<Vec<Notification>>>,
|
||||
/// Notification channels
|
||||
channels: Arc<RwLock<Vec<NotificationChannel>>>,
|
||||
/// Maximum notifications to keep in history
|
||||
max_history: usize,
|
||||
/// Whether the manager is enabled
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl AlertManager {
|
||||
/// Create a new alert manager
|
||||
pub async fn new() -> MonitoringResult<Self> {
|
||||
Ok(Self {
|
||||
rules: Arc::new(RwLock::new(HashMap::new())),
|
||||
active_alerts: Arc::new(RwLock::new(HashMap::new())),
|
||||
notifications: Arc::new(RwLock::new(Vec::new())),
|
||||
channels: Arc::new(RwLock::new(Vec::new())),
|
||||
max_history: 10000,
|
||||
enabled: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Register an alert rule
|
||||
pub fn register_rule(&self, rule: AlertRule) {
|
||||
self.rules.write().insert(rule.name.clone(), rule);
|
||||
}
|
||||
|
||||
/// Remove an alert rule
|
||||
pub fn remove_rule(&self, name: &str) {
|
||||
self.rules.write().remove(name);
|
||||
self.active_alerts.write().remove(name);
|
||||
}
|
||||
|
||||
/// Add a notification channel
|
||||
pub fn add_channel(&self, channel: NotificationChannel) {
|
||||
self.channels.write().push(channel);
|
||||
}
|
||||
|
||||
/// Evaluate a metric value against all rules
|
||||
pub async fn evaluate(&self, metric: &str, value: Option<f64>) -> MonitoringResult<()> {
|
||||
let rules = self.rules.read();
|
||||
let now = Utc::now();
|
||||
|
||||
for rule in rules.values() {
|
||||
if rule.metric != metric {
|
||||
continue;
|
||||
}
|
||||
|
||||
let condition_met = rule.condition.evaluate(value);
|
||||
let mut active_alerts = self.active_alerts.write();
|
||||
|
||||
if condition_met {
|
||||
if let Some(active) = active_alerts.get_mut(&rule.name) {
|
||||
// Already tracking this alert
|
||||
active.last_evaluation = now;
|
||||
active.value = value;
|
||||
|
||||
// Check if duration has passed
|
||||
if active.state == AlertState::Pending {
|
||||
let elapsed = now - active.first_firing;
|
||||
if elapsed >= Duration::from_std(rule.duration).unwrap_or(Duration::zero())
|
||||
{
|
||||
active.state = AlertState::Firing;
|
||||
self.fire_alert(rule, value).await?;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// New alert
|
||||
let state = if rule.duration.is_zero() {
|
||||
AlertState::Firing
|
||||
} else {
|
||||
AlertState::Pending
|
||||
};
|
||||
|
||||
active_alerts.insert(
|
||||
rule.name.clone(),
|
||||
ActiveAlert {
|
||||
rule: rule.clone(),
|
||||
first_firing: now,
|
||||
last_evaluation: now,
|
||||
state,
|
||||
value,
|
||||
},
|
||||
);
|
||||
|
||||
if state == AlertState::Firing {
|
||||
drop(active_alerts);
|
||||
self.fire_alert(rule, value).await?;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Condition no longer met
|
||||
if let Some(active) = active_alerts.remove(&rule.name)
|
||||
&& active.state == AlertState::Firing
|
||||
{
|
||||
// Send resolved notification
|
||||
drop(active_alerts);
|
||||
self.resolve_alert(rule, value).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fire an alert
|
||||
async fn fire_alert(&self, rule: &AlertRule, value: Option<f64>) -> MonitoringResult<()> {
|
||||
let notification = Notification::from_rule(rule, AlertState::Firing, value);
|
||||
|
||||
// Add to history
|
||||
{
|
||||
let mut notifications = self.notifications.write();
|
||||
if notifications.len() >= self.max_history {
|
||||
notifications.remove(0);
|
||||
}
|
||||
notifications.push(notification.clone());
|
||||
}
|
||||
|
||||
// Dispatch to channels
|
||||
self.dispatch_notification(¬ification).await?;
|
||||
|
||||
tracing::warn!(
|
||||
rule = %rule.name,
|
||||
severity = %rule.severity.as_str(),
|
||||
value = ?value,
|
||||
"Alert fired: {}",
|
||||
rule.description
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve an alert
|
||||
async fn resolve_alert(&self, rule: &AlertRule, value: Option<f64>) -> MonitoringResult<()> {
|
||||
let notification = Notification::from_rule(rule, AlertState::Resolved, value);
|
||||
|
||||
// Add to history
|
||||
{
|
||||
let mut notifications = self.notifications.write();
|
||||
if notifications.len() >= self.max_history {
|
||||
notifications.remove(0);
|
||||
}
|
||||
notifications.push(notification.clone());
|
||||
}
|
||||
|
||||
// Dispatch to channels
|
||||
self.dispatch_notification(¬ification).await?;
|
||||
|
||||
tracing::info!(
|
||||
rule = %rule.name,
|
||||
"Alert resolved: {}",
|
||||
rule.description
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Dispatch notification to all configured channels
|
||||
async fn dispatch_notification(&self, notification: &Notification) -> MonitoringResult<()> {
|
||||
let channels = self.channels.read();
|
||||
|
||||
for channel in channels.iter() {
|
||||
// Check severity filter
|
||||
let should_send = match (channel.min_severity, notification.severity) {
|
||||
(AlertSeverity::Critical, AlertSeverity::Critical) => true,
|
||||
(AlertSeverity::Warning, AlertSeverity::Critical | AlertSeverity::Warning) => true,
|
||||
(AlertSeverity::Info, _) => true,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if !should_send {
|
||||
continue;
|
||||
}
|
||||
|
||||
match channel.channel_type {
|
||||
ChannelType::Console => {
|
||||
tracing::info!(
|
||||
channel = %channel.name,
|
||||
severity = %notification.severity.as_str(),
|
||||
state = ?notification.state,
|
||||
"Alert: {} - {}",
|
||||
notification.rule_name,
|
||||
notification.message
|
||||
);
|
||||
}
|
||||
ChannelType::Webhook => {
|
||||
// In a real implementation, this would make an HTTP POST
|
||||
if let Some(url) = channel.config.get("url") {
|
||||
tracing::debug!("Would send webhook to: {}", url);
|
||||
}
|
||||
}
|
||||
ChannelType::Slack => {
|
||||
// In a real implementation, this would send to Slack
|
||||
if let Some(webhook) = channel.config.get("webhook_url") {
|
||||
tracing::debug!("Would send Slack notification to: {}", webhook);
|
||||
}
|
||||
}
|
||||
ChannelType::PagerDuty => {
|
||||
// In a real implementation, this would trigger PagerDuty
|
||||
if let Some(key) = channel.config.get("routing_key") {
|
||||
tracing::debug!("Would trigger PagerDuty with key: {}", key);
|
||||
}
|
||||
}
|
||||
ChannelType::Email => {
|
||||
// Email would be sent via lettre if the feature is enabled
|
||||
if let Some(to) = channel.config.get("to") {
|
||||
tracing::debug!("Would send email to: {}", to);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get all active alerts
|
||||
pub async fn active_alerts(&self) -> MonitoringResult<Vec<Notification>> {
|
||||
let active = self.active_alerts.read();
|
||||
let notifications: Vec<Notification> = active
|
||||
.values()
|
||||
.filter(|a| a.state == AlertState::Firing)
|
||||
.map(|a| Notification::from_rule(&a.rule, a.state, a.value))
|
||||
.collect();
|
||||
Ok(notifications)
|
||||
}
|
||||
|
||||
/// Get notification history
|
||||
pub fn notification_history(&self, limit: usize) -> Vec<Notification> {
|
||||
let notifications = self.notifications.read();
|
||||
notifications.iter().rev().take(limit).cloned().collect()
|
||||
}
|
||||
|
||||
/// Get alert counts by severity
|
||||
pub fn alert_counts(&self) -> HashMap<AlertSeverity, usize> {
|
||||
let active = self.active_alerts.read();
|
||||
let mut counts = HashMap::new();
|
||||
|
||||
for alert in active.values() {
|
||||
if alert.state == AlertState::Firing {
|
||||
*counts.entry(alert.rule.severity).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
counts
|
||||
}
|
||||
|
||||
/// Check if the alert manager is enabled
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
/// Get registered rule count
|
||||
pub fn rule_count(&self) -> usize {
|
||||
self.rules.read().len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-defined alert rules for common ML scenarios
|
||||
pub mod presets {
|
||||
use super::{AlertRule, AlertCondition, AlertSeverity};
|
||||
|
||||
/// High GPU memory usage alert
|
||||
pub fn high_gpu_memory(threshold_percent: f64) -> AlertRule {
|
||||
AlertRule::new(
|
||||
"high_gpu_memory",
|
||||
"rtx_gpu_memory_percent",
|
||||
AlertCondition::GreaterThan(threshold_percent),
|
||||
AlertSeverity::Warning,
|
||||
)
|
||||
.with_description(&format!("GPU memory usage exceeds {threshold_percent}%"))
|
||||
.with_duration(std::time::Duration::from_secs(60))
|
||||
}
|
||||
|
||||
/// High inference latency alert
|
||||
pub fn high_inference_latency(threshold_secs: f64) -> AlertRule {
|
||||
AlertRule::new(
|
||||
"high_inference_latency",
|
||||
"rtx_inference_latency_p99",
|
||||
AlertCondition::GreaterThan(threshold_secs),
|
||||
AlertSeverity::Warning,
|
||||
)
|
||||
.with_description(&format!("P99 inference latency exceeds {threshold_secs}s"))
|
||||
.with_duration(std::time::Duration::from_secs(300))
|
||||
}
|
||||
|
||||
/// High error rate alert
|
||||
pub fn high_error_rate(threshold_percent: f64) -> AlertRule {
|
||||
AlertRule::new(
|
||||
"high_error_rate",
|
||||
"rtx_inference_error_rate",
|
||||
AlertCondition::GreaterThan(threshold_percent),
|
||||
AlertSeverity::Critical,
|
||||
)
|
||||
.with_description(&format!("Error rate exceeds {threshold_percent}%"))
|
||||
.with_duration(std::time::Duration::from_secs(60))
|
||||
}
|
||||
|
||||
/// Model not loaded alert
|
||||
pub fn model_not_loaded() -> AlertRule {
|
||||
AlertRule::new(
|
||||
"model_not_loaded",
|
||||
"rtx_model_loaded",
|
||||
AlertCondition::Equal(0.0),
|
||||
AlertSeverity::Critical,
|
||||
)
|
||||
.with_description("No model is loaded")
|
||||
.with_duration(std::time::Duration::from_secs(0))
|
||||
}
|
||||
|
||||
/// Queue depth high alert
|
||||
pub fn queue_depth_high(threshold: f64) -> AlertRule {
|
||||
AlertRule::new(
|
||||
"queue_depth_high",
|
||||
"rtx_queue_depth",
|
||||
AlertCondition::GreaterThan(threshold),
|
||||
AlertSeverity::Warning,
|
||||
)
|
||||
.with_description(&format!("Request queue depth exceeds {threshold}"))
|
||||
.with_duration(std::time::Duration::from_secs(120))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_alert_condition_evaluate() {
|
||||
assert!(AlertCondition::GreaterThan(50.0).evaluate(Some(60.0)));
|
||||
assert!(!AlertCondition::GreaterThan(50.0).evaluate(Some(40.0)));
|
||||
assert!(AlertCondition::LessThan(50.0).evaluate(Some(40.0)));
|
||||
assert!(AlertCondition::Absent.evaluate(None));
|
||||
assert!(!AlertCondition::Absent.evaluate(Some(1.0)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_alert_rule_builder() {
|
||||
let rule = AlertRule::new(
|
||||
"test_rule",
|
||||
"test_metric",
|
||||
AlertCondition::GreaterThan(90.0),
|
||||
AlertSeverity::Critical,
|
||||
)
|
||||
.with_description("Test description")
|
||||
.with_duration(std::time::Duration::from_secs(60))
|
||||
.with_label("env", "prod")
|
||||
.with_annotation("runbook", "https://docs.example.com/runbook");
|
||||
|
||||
assert_eq!(rule.name, "test_rule");
|
||||
assert_eq!(rule.severity, AlertSeverity::Critical);
|
||||
assert_eq!(rule.labels.get("env"), Some(&"prod".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_alert_manager() {
|
||||
let manager = AlertManager::new().await.unwrap();
|
||||
|
||||
// Register a rule
|
||||
let rule = AlertRule::new(
|
||||
"test_alert",
|
||||
"test_metric",
|
||||
AlertCondition::GreaterThan(80.0),
|
||||
AlertSeverity::Warning,
|
||||
);
|
||||
manager.register_rule(rule);
|
||||
|
||||
assert_eq!(manager.rule_count(), 1);
|
||||
|
||||
// Evaluate - should not fire
|
||||
manager.evaluate("test_metric", Some(50.0)).await.unwrap();
|
||||
let active = manager.active_alerts().await.unwrap();
|
||||
assert!(active.is_empty());
|
||||
|
||||
// Evaluate - should fire
|
||||
manager.evaluate("test_metric", Some(90.0)).await.unwrap();
|
||||
let active = manager.active_alerts().await.unwrap();
|
||||
assert_eq!(active.len(), 1);
|
||||
|
||||
// Evaluate - should resolve
|
||||
manager.evaluate("test_metric", Some(50.0)).await.unwrap();
|
||||
let active = manager.active_alerts().await.unwrap();
|
||||
assert!(active.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_preset_rules() {
|
||||
let rule = presets::high_gpu_memory(90.0);
|
||||
assert_eq!(rule.name, "high_gpu_memory");
|
||||
assert_eq!(rule.severity, AlertSeverity::Warning);
|
||||
|
||||
let rule = presets::high_error_rate(5.0);
|
||||
assert_eq!(rule.severity, AlertSeverity::Critical);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_notification_from_rule() {
|
||||
let rule = AlertRule::new(
|
||||
"test",
|
||||
"metric",
|
||||
AlertCondition::GreaterThan(1.0),
|
||||
AlertSeverity::Warning,
|
||||
)
|
||||
.with_description("Test alert");
|
||||
|
||||
let notification = Notification::from_rule(&rule, AlertState::Firing, Some(2.0));
|
||||
assert_eq!(notification.rule_name, "test");
|
||||
assert_eq!(notification.state, AlertState::Firing);
|
||||
assert_eq!(notification.value, Some(2.0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
//! Metrics collector implementation.
|
||||
|
||||
use crate::MonitoringResult;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
|
||||
/// Configuration for metrics collector.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CollectorConfig {
|
||||
pub enable_prometheus: bool,
|
||||
pub collection_interval: Duration,
|
||||
pub enable_system_metrics: bool,
|
||||
}
|
||||
|
||||
impl Default for CollectorConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enable_prometheus: true,
|
||||
collection_interval: Duration::from_secs(30),
|
||||
enable_system_metrics: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// System metrics data.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SystemMetrics {
|
||||
pub cpu_usage: f64,
|
||||
pub memory_usage: f64,
|
||||
pub disk_usage: f64,
|
||||
pub network_rx_bytes: u64,
|
||||
pub network_tx_bytes: u64,
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
impl Default for SystemMetrics {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
cpu_usage: 0.0,
|
||||
memory_usage: 0.0,
|
||||
disk_usage: 0.0,
|
||||
network_rx_bytes: 0,
|
||||
network_tx_bytes: 0,
|
||||
timestamp: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metrics collector for gathering system and application metrics.
|
||||
#[derive(Debug)]
|
||||
pub struct MetricsCollector {
|
||||
config: CollectorConfig,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl MetricsCollector {
|
||||
pub async fn new(config: CollectorConfig) -> MonitoringResult<Self> {
|
||||
Ok(Self {
|
||||
config,
|
||||
enabled: true,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
pub async fn collect_system_metrics(&self) -> MonitoringResult<SystemMetrics> {
|
||||
// Placeholder implementation
|
||||
Ok(SystemMetrics {
|
||||
cpu_usage: 25.5,
|
||||
memory_usage: 60.2,
|
||||
disk_usage: 45.8,
|
||||
network_rx_bytes: 1024 * 1024,
|
||||
network_tx_bytes: 512 * 1024,
|
||||
timestamp: chrono::Utc::now(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn export_prometheus(&self) -> MonitoringResult<String> {
|
||||
// Placeholder Prometheus metrics export
|
||||
let metrics = self.collect_system_metrics().await?;
|
||||
|
||||
Ok(format!(
|
||||
"# HELP cpu_usage_percent CPU usage percentage\n\
|
||||
# TYPE cpu_usage_percent gauge\n\
|
||||
cpu_usage_percent {}\n\
|
||||
# HELP memory_usage_percent Memory usage percentage\n\
|
||||
# TYPE memory_usage_percent gauge\n\
|
||||
memory_usage_percent {}\n\
|
||||
# HELP disk_usage_percent Disk usage percentage\n\
|
||||
# TYPE disk_usage_percent gauge\n\
|
||||
disk_usage_percent {}\n",
|
||||
metrics.cpu_usage, metrics.memory_usage, metrics.disk_usage
|
||||
))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! Error types for the RTX Monitoring system.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Result type for RTX Monitoring operations.
|
||||
pub type MonitoringResult<T> = Result<T, MonitoringError>;
|
||||
|
||||
/// Comprehensive error types for the monitoring system.
|
||||
#[derive(Error, Debug)]
|
||||
pub enum MonitoringError {
|
||||
/// Metrics collection error.
|
||||
#[error("Metrics collection error: {details}")]
|
||||
MetricsError { details: String },
|
||||
|
||||
/// Health check error.
|
||||
#[error("Health check error: {check_name} - {details}")]
|
||||
HealthCheckError { check_name: String, details: String },
|
||||
|
||||
/// Alert management error.
|
||||
#[error("Alert error: {details}")]
|
||||
AlertError { details: String },
|
||||
|
||||
/// Telemetry error.
|
||||
#[error("Telemetry error: {details}")]
|
||||
TelemetryError { details: String },
|
||||
|
||||
/// Configuration error.
|
||||
#[error("Configuration error: {details}")]
|
||||
ConfigError { details: String },
|
||||
|
||||
/// Server error.
|
||||
#[error("Server error: {details}")]
|
||||
ServerError { details: String },
|
||||
|
||||
/// I/O error.
|
||||
#[error("I/O error: {source}")]
|
||||
IoError {
|
||||
#[from]
|
||||
source: std::io::Error,
|
||||
},
|
||||
|
||||
/// Serialization error.
|
||||
#[error("Serialization error: {source}")]
|
||||
SerializationError {
|
||||
#[from]
|
||||
source: serde_json::Error,
|
||||
},
|
||||
}
|
||||
|
||||
impl MonitoringError {
|
||||
/// Create a metrics error.
|
||||
pub fn metrics_error(details: impl Into<String>) -> Self {
|
||||
Self::MetricsError {
|
||||
details: details.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a health check error.
|
||||
pub fn health_check_error(check_name: impl Into<String>, details: impl Into<String>) -> Self {
|
||||
Self::HealthCheckError {
|
||||
check_name: check_name.into(),
|
||||
details: details.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
//! Health check management.
|
||||
|
||||
use crate::MonitoringResult;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Health check status.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum HealthStatus {
|
||||
Healthy,
|
||||
Degraded,
|
||||
Unhealthy,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// Health check trait.
|
||||
#[async_trait]
|
||||
pub trait HealthCheck: Send + Sync {
|
||||
async fn check(&self) -> MonitoringResult<HealthStatus>;
|
||||
fn name(&self) -> &str;
|
||||
}
|
||||
|
||||
/// Health manager for coordinating health checks.
|
||||
pub struct HealthManager {
|
||||
checks: Vec<Box<dyn HealthCheck>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HealthManager {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("HealthManager")
|
||||
.field("checks_count", &self.checks.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl HealthManager {
|
||||
pub async fn new() -> MonitoringResult<Self> {
|
||||
Ok(Self { checks: Vec::new() })
|
||||
}
|
||||
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
pub async fn run_all_checks(&self) -> MonitoringResult<()> {
|
||||
for check in &self.checks {
|
||||
match check.check().await {
|
||||
Ok(status) => {
|
||||
tracing::debug!("Health check '{}': {:?}", check.name(), status);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Health check '{}' failed: {}", check.name(), e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn overall_status(&self) -> MonitoringResult<HealthStatus> {
|
||||
Ok(HealthStatus::Healthy)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
//! # RTX Monitoring - Observability and Metrics Collection
|
||||
//!
|
||||
//! RTX Monitoring provides comprehensive observability for RustyTorch production
|
||||
//! deployments, including metrics collection, distributed tracing, health checks,
|
||||
//! and alerting.
|
||||
//!
|
||||
//! ## Core Features
|
||||
//!
|
||||
//! - Prometheus metrics collection and export
|
||||
//! - OpenTelemetry distributed tracing
|
||||
//! - System and application health monitoring
|
||||
//! - Custom metrics and dashboards
|
||||
//! - Alert management and notification
|
||||
//! - Performance monitoring and SLI/SLO tracking
|
||||
|
||||
#![deny(clippy::unwrap_used)]
|
||||
#![cfg_attr(test, allow(clippy::unwrap_used))]
|
||||
|
||||
pub mod alerts;
|
||||
pub mod collector;
|
||||
pub mod error;
|
||||
pub mod health;
|
||||
pub mod metrics;
|
||||
pub mod telemetry;
|
||||
pub mod trace_utils;
|
||||
|
||||
pub use alerts::{
|
||||
AlertCondition, AlertManager, AlertRule, AlertSeverity, AlertState, ChannelType, Notification,
|
||||
NotificationChannel, presets as alert_presets,
|
||||
};
|
||||
pub use collector::{CollectorConfig, MetricsCollector, SystemMetrics};
|
||||
pub use error::{MonitoringError, MonitoringResult};
|
||||
pub use health::{HealthCheck, HealthManager, HealthStatus};
|
||||
pub use metrics::{
|
||||
BATCH_SIZE_BUCKETS, CustomMetric, INFERENCE_LATENCY_BUCKETS, InferenceMetrics, MetricRegistry,
|
||||
MetricType, TrainingMetrics,
|
||||
};
|
||||
pub use telemetry::{
|
||||
SpanContext, SpanData, SpanEvent, SpanGuard, SpanStatus, TelemetryManager, TraceConfig,
|
||||
generate_span_id, generate_trace_id,
|
||||
};
|
||||
|
||||
use std::sync::Arc;
|
||||
// Import tracing macros with alias to avoid conflict with local tracing module
|
||||
|
||||
/// Shared monitoring system instance
|
||||
pub type SharedMonitoringSystem = Arc<MonitoringSystem>;
|
||||
|
||||
/// Main monitoring system that coordinates all observability components
|
||||
#[derive(Debug)]
|
||||
pub struct MonitoringSystem {
|
||||
/// Metrics collector
|
||||
pub metrics: Arc<MetricsCollector>,
|
||||
/// Health manager
|
||||
pub health: Arc<HealthManager>,
|
||||
/// Alert manager
|
||||
pub alerts: Arc<AlertManager>,
|
||||
/// Telemetry manager
|
||||
pub telemetry: Arc<TelemetryManager>,
|
||||
/// System configuration
|
||||
config: MonitoringConfig,
|
||||
}
|
||||
|
||||
/// Configuration for the monitoring system
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MonitoringConfig {
|
||||
/// Enable Prometheus metrics
|
||||
pub enable_prometheus: bool,
|
||||
/// Enable OpenTelemetry tracing
|
||||
pub enable_telemetry: bool,
|
||||
/// Enable health checks
|
||||
pub enable_health_checks: bool,
|
||||
/// Enable alerting
|
||||
pub enable_alerting: bool,
|
||||
/// Metrics collection interval
|
||||
pub collection_interval: std::time::Duration,
|
||||
/// Health check interval
|
||||
pub health_check_interval: std::time::Duration,
|
||||
/// Metrics server port
|
||||
pub metrics_port: u16,
|
||||
/// Health check port
|
||||
pub health_port: u16,
|
||||
}
|
||||
|
||||
impl Default for MonitoringConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enable_prometheus: true,
|
||||
enable_telemetry: true,
|
||||
enable_health_checks: true,
|
||||
enable_alerting: false,
|
||||
collection_interval: std::time::Duration::from_secs(30),
|
||||
health_check_interval: std::time::Duration::from_secs(10),
|
||||
metrics_port: 9090,
|
||||
health_port: 8080,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MonitoringSystem {
|
||||
/// Create a new monitoring system
|
||||
pub async fn new(config: MonitoringConfig) -> MonitoringResult<Self> {
|
||||
let metrics = Arc::new(
|
||||
MetricsCollector::new(CollectorConfig {
|
||||
enable_prometheus: config.enable_prometheus,
|
||||
collection_interval: config.collection_interval,
|
||||
..Default::default()
|
||||
})
|
||||
.await?,
|
||||
);
|
||||
|
||||
let health = Arc::new(HealthManager::new().await?);
|
||||
let alerts = Arc::new(AlertManager::new().await?);
|
||||
let telemetry = Arc::new(TelemetryManager::new(TraceConfig::default()).await?);
|
||||
|
||||
let enable_prometheus = config.enable_prometheus;
|
||||
let enable_health_checks = config.enable_health_checks;
|
||||
|
||||
let system = Self {
|
||||
metrics,
|
||||
health,
|
||||
alerts,
|
||||
telemetry,
|
||||
config,
|
||||
};
|
||||
|
||||
// Start background tasks
|
||||
if enable_prometheus {
|
||||
system.start_metrics_server().await?;
|
||||
}
|
||||
|
||||
if enable_health_checks {
|
||||
system.start_health_checks().await?;
|
||||
}
|
||||
|
||||
Ok(system)
|
||||
}
|
||||
|
||||
/// Start the metrics HTTP server
|
||||
async fn start_metrics_server(&self) -> MonitoringResult<()> {
|
||||
let metrics = self.metrics.clone();
|
||||
let port = self.config.metrics_port;
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = run_metrics_server(metrics, port).await {
|
||||
tracing::error!("Metrics server error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Start health check background tasks
|
||||
async fn start_health_checks(&self) -> MonitoringResult<()> {
|
||||
let health = self.health.clone();
|
||||
let interval = self.config.health_check_interval;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut ticker = tokio::time::interval(interval);
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
if let Err(e) = health.run_all_checks().await {
|
||||
tracing::error!("Health check error: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get system metrics summary
|
||||
pub async fn get_metrics_summary(&self) -> MonitoringResult<MetricsSummary> {
|
||||
let system_metrics = self.metrics.collect_system_metrics().await?;
|
||||
let health_status = self.health.overall_status().await?;
|
||||
let active_alerts = self.alerts.active_alerts().await?;
|
||||
|
||||
Ok(MetricsSummary {
|
||||
system_metrics,
|
||||
health_status,
|
||||
active_alerts: active_alerts.len(),
|
||||
uptime: self.get_uptime(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get system uptime
|
||||
fn get_uptime(&self) -> std::time::Duration {
|
||||
// Placeholder implementation
|
||||
std::time::Duration::from_secs(3600)
|
||||
}
|
||||
}
|
||||
|
||||
/// Summary of system metrics
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct MetricsSummary {
|
||||
pub system_metrics: SystemMetrics,
|
||||
pub health_status: HealthStatus,
|
||||
pub active_alerts: usize,
|
||||
pub uptime: std::time::Duration,
|
||||
}
|
||||
|
||||
/// Run the metrics HTTP server
|
||||
async fn run_metrics_server(metrics: Arc<MetricsCollector>, port: u16) -> MonitoringResult<()> {
|
||||
use axum::{Router, routing::get};
|
||||
|
||||
let app = Router::new().route(
|
||||
"/metrics",
|
||||
get(move || async move {
|
||||
match metrics.export_prometheus().await {
|
||||
Ok(metrics_text) => axum::response::Response::builder()
|
||||
.header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
|
||||
.body(metrics_text)
|
||||
.expect("failed to build successful metrics response"),
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to export Prometheus metrics: {}", e);
|
||||
axum::response::Response::builder()
|
||||
.status(500)
|
||||
.body("Internal Server Error".to_string())
|
||||
.expect("failed to build error response")
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let addr = format!("0.0.0.0:{port}");
|
||||
let listener =
|
||||
tokio::net::TcpListener::bind(&addr)
|
||||
.await
|
||||
.map_err(|e| MonitoringError::ConfigError {
|
||||
details: format!("Failed to bind metrics server to {addr}: {e}"),
|
||||
})?;
|
||||
|
||||
tracing::info!("Metrics server listening on {}", addr);
|
||||
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.map_err(|e| MonitoringError::ServerError {
|
||||
details: e.to_string(),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create a production-ready monitoring system
|
||||
pub async fn create_production_monitoring() -> MonitoringResult<MonitoringSystem> {
|
||||
let config = MonitoringConfig {
|
||||
enable_prometheus: true,
|
||||
enable_telemetry: true,
|
||||
enable_health_checks: true,
|
||||
enable_alerting: true,
|
||||
collection_interval: std::time::Duration::from_secs(15),
|
||||
health_check_interval: std::time::Duration::from_secs(5),
|
||||
metrics_port: 9090,
|
||||
health_port: 8080,
|
||||
};
|
||||
|
||||
MonitoringSystem::new(config).await
|
||||
}
|
||||
|
||||
/// Create a development monitoring system
|
||||
pub async fn create_development_monitoring() -> MonitoringResult<MonitoringSystem> {
|
||||
let config = MonitoringConfig {
|
||||
enable_prometheus: true,
|
||||
enable_telemetry: false,
|
||||
enable_health_checks: true,
|
||||
enable_alerting: false,
|
||||
collection_interval: std::time::Duration::from_secs(60),
|
||||
health_check_interval: std::time::Duration::from_secs(30),
|
||||
metrics_port: 9090,
|
||||
health_port: 8080,
|
||||
};
|
||||
|
||||
MonitoringSystem::new(config).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_monitoring_system_creation() {
|
||||
let config = MonitoringConfig::default();
|
||||
let system = MonitoringSystem::new(config).await.unwrap();
|
||||
|
||||
assert!(system.metrics.is_enabled());
|
||||
assert!(system.health.is_enabled());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_production_monitoring_setup() {
|
||||
let system = create_production_monitoring().await.unwrap();
|
||||
assert!(system.config.enable_prometheus);
|
||||
assert!(system.config.enable_telemetry);
|
||||
assert!(system.config.enable_health_checks);
|
||||
assert!(system.config.enable_alerting);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_development_monitoring_setup() {
|
||||
let system = create_development_monitoring().await.unwrap();
|
||||
assert!(system.config.enable_prometheus);
|
||||
assert!(!system.config.enable_telemetry);
|
||||
assert!(system.config.enable_health_checks);
|
||||
assert!(!system.config.enable_alerting);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,546 @@
|
||||
//! Telemetry and distributed tracing.
|
||||
//!
|
||||
//! This module provides distributed tracing capabilities for RustyTorch++,
|
||||
//! including trace context propagation, span instrumentation, and export
|
||||
//! to various tracing backends.
|
||||
|
||||
use crate::MonitoringResult;
|
||||
use parking_lot::RwLock;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use tracing::info;
|
||||
|
||||
/// Trace ID generator using atomic counter + timestamp for uniqueness
|
||||
static TRACE_COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// Generate a unique trace ID
|
||||
pub fn generate_trace_id() -> String {
|
||||
let counter = TRACE_COUNTER.fetch_add(1, Ordering::SeqCst);
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
format!("{:016x}{:016x}", timestamp as u64, counter)
|
||||
}
|
||||
|
||||
/// Generate a unique span ID
|
||||
pub fn generate_span_id() -> String {
|
||||
let counter = TRACE_COUNTER.fetch_add(1, Ordering::SeqCst);
|
||||
format!("{counter:016x}")
|
||||
}
|
||||
|
||||
/// Span context for distributed tracing.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SpanContext {
|
||||
/// Unique trace identifier (propagated across services)
|
||||
pub trace_id: String,
|
||||
/// Unique span identifier
|
||||
pub span_id: String,
|
||||
/// Parent span ID (if any)
|
||||
pub parent_span_id: Option<String>,
|
||||
/// Sampling decision
|
||||
pub sampled: bool,
|
||||
/// Baggage items (propagated context)
|
||||
#[serde(default)]
|
||||
pub baggage: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl SpanContext {
|
||||
/// Create a new root span context
|
||||
pub fn new_root() -> Self {
|
||||
Self {
|
||||
trace_id: generate_trace_id(),
|
||||
span_id: generate_span_id(),
|
||||
parent_span_id: None,
|
||||
sampled: true,
|
||||
baggage: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a child span context
|
||||
pub fn child(&self) -> Self {
|
||||
Self {
|
||||
trace_id: self.trace_id.clone(),
|
||||
span_id: generate_span_id(),
|
||||
parent_span_id: Some(self.span_id.clone()),
|
||||
sampled: self.sampled,
|
||||
baggage: self.baggage.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add baggage item
|
||||
pub fn with_baggage(mut self, key: &str, value: &str) -> Self {
|
||||
self.baggage.insert(key.to_string(), value.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Serialize to W3C Trace Context format for HTTP headers
|
||||
pub fn to_traceparent(&self) -> String {
|
||||
let flags = if self.sampled { "01" } else { "00" };
|
||||
format!("00-{}-{}-{}", self.trace_id, self.span_id, flags)
|
||||
}
|
||||
|
||||
/// Parse from W3C Trace Context format
|
||||
pub fn from_traceparent(header: &str) -> Option<Self> {
|
||||
let parts: Vec<&str> = header.split('-').collect();
|
||||
if parts.len() != 4 || parts[0] != "00" {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(Self {
|
||||
trace_id: parts[1].to_string(),
|
||||
span_id: generate_span_id(), // New span ID for this service
|
||||
parent_span_id: Some(parts[2].to_string()),
|
||||
sampled: parts[3] == "01",
|
||||
baggage: HashMap::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Trace configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TraceConfig {
|
||||
/// Service name for tracing
|
||||
pub service_name: String,
|
||||
/// Sample rate (0.0 to 1.0)
|
||||
pub sample_rate: f64,
|
||||
/// Enable console output
|
||||
pub enable_console: bool,
|
||||
/// Enable JSON output
|
||||
pub enable_json: bool,
|
||||
/// Jaeger endpoint (if enabled)
|
||||
pub jaeger_endpoint: Option<String>,
|
||||
/// OTLP endpoint (if enabled)
|
||||
pub otlp_endpoint: Option<String>,
|
||||
/// Maximum spans to buffer
|
||||
pub max_spans_buffer: usize,
|
||||
/// Batch export interval
|
||||
pub export_interval: std::time::Duration,
|
||||
}
|
||||
|
||||
impl Default for TraceConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
service_name: "rtx-service".to_string(),
|
||||
sample_rate: 1.0,
|
||||
enable_console: true,
|
||||
enable_json: false,
|
||||
jaeger_endpoint: None,
|
||||
otlp_endpoint: None,
|
||||
max_spans_buffer: 10000,
|
||||
export_interval: std::time::Duration::from_secs(5),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TraceConfig {
|
||||
/// Create a production configuration
|
||||
pub fn production(service_name: &str) -> Self {
|
||||
Self {
|
||||
service_name: service_name.to_string(),
|
||||
sample_rate: 0.1, // Sample 10% in production
|
||||
enable_console: false,
|
||||
enable_json: true,
|
||||
jaeger_endpoint: None,
|
||||
otlp_endpoint: None,
|
||||
max_spans_buffer: 50000,
|
||||
export_interval: std::time::Duration::from_secs(10),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a development configuration
|
||||
pub fn development(service_name: &str) -> Self {
|
||||
Self {
|
||||
service_name: service_name.to_string(),
|
||||
sample_rate: 1.0, // Sample everything in dev
|
||||
enable_console: true,
|
||||
enable_json: false,
|
||||
jaeger_endpoint: None,
|
||||
otlp_endpoint: None,
|
||||
max_spans_buffer: 1000,
|
||||
export_interval: std::time::Duration::from_secs(1),
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable Jaeger export
|
||||
pub fn with_jaeger(mut self, endpoint: &str) -> Self {
|
||||
self.jaeger_endpoint = Some(endpoint.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable OTLP export
|
||||
pub fn with_otlp(mut self, endpoint: &str) -> Self {
|
||||
self.otlp_endpoint = Some(endpoint.to_string());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Recorded span data for export
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpanData {
|
||||
/// Span context
|
||||
pub context: SpanContext,
|
||||
/// Operation name
|
||||
pub operation_name: String,
|
||||
/// Start timestamp
|
||||
pub start_time: chrono::DateTime<chrono::Utc>,
|
||||
/// End timestamp (if completed)
|
||||
pub end_time: Option<chrono::DateTime<chrono::Utc>>,
|
||||
/// Duration in microseconds
|
||||
pub duration_us: Option<u64>,
|
||||
/// Span status
|
||||
pub status: SpanStatus,
|
||||
/// Span attributes
|
||||
pub attributes: HashMap<String, String>,
|
||||
/// Span events
|
||||
pub events: Vec<SpanEvent>,
|
||||
}
|
||||
|
||||
/// Span status
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub enum SpanStatus {
|
||||
/// Unset status
|
||||
#[default]
|
||||
Unset,
|
||||
/// Operation completed successfully
|
||||
Ok,
|
||||
/// Operation failed
|
||||
Error(String),
|
||||
}
|
||||
|
||||
/// Span event (annotation)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpanEvent {
|
||||
/// Event name
|
||||
pub name: String,
|
||||
/// Event timestamp
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
/// Event attributes
|
||||
pub attributes: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Telemetry manager for distributed tracing.
|
||||
#[derive(Debug)]
|
||||
pub struct TelemetryManager {
|
||||
config: TraceConfig,
|
||||
/// Buffer for completed spans
|
||||
spans_buffer: Arc<RwLock<Vec<SpanData>>>,
|
||||
/// Active spans by trace ID
|
||||
active_spans: Arc<RwLock<HashMap<String, SpanData>>>,
|
||||
/// Whether the manager is enabled
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl TelemetryManager {
|
||||
/// Create a new telemetry manager
|
||||
pub async fn new(config: TraceConfig) -> MonitoringResult<Self> {
|
||||
let manager = Self {
|
||||
config: config.clone(),
|
||||
spans_buffer: Arc::new(RwLock::new(Vec::new())),
|
||||
active_spans: Arc::new(RwLock::new(HashMap::new())),
|
||||
enabled: true,
|
||||
};
|
||||
|
||||
// Initialize tracing subscriber
|
||||
manager.init_tracing()?;
|
||||
|
||||
info!(
|
||||
service = %config.service_name,
|
||||
sample_rate = %config.sample_rate,
|
||||
"Telemetry manager initialized"
|
||||
);
|
||||
|
||||
Ok(manager)
|
||||
}
|
||||
|
||||
/// Initialize the tracing subscriber
|
||||
fn init_tracing(&self) -> MonitoringResult<()> {
|
||||
use tracing_subscriber::{EnvFilter, fmt, prelude::*};
|
||||
|
||||
let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
|
||||
|
||||
let subscriber = tracing_subscriber::registry().with(filter);
|
||||
|
||||
if self.config.enable_json {
|
||||
let json_layer = fmt::layer()
|
||||
.json()
|
||||
.with_target(true)
|
||||
.with_thread_ids(true)
|
||||
.with_file(true)
|
||||
.with_line_number(true);
|
||||
|
||||
let _ = subscriber.with(json_layer).try_init();
|
||||
} else if self.config.enable_console {
|
||||
let fmt_layer = fmt::layer()
|
||||
.with_target(true)
|
||||
.with_thread_ids(false)
|
||||
.with_file(false)
|
||||
.with_line_number(false);
|
||||
|
||||
let _ = subscriber.with(fmt_layer).try_init();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Start a new span
|
||||
pub fn start_span(&self, operation_name: &str, parent: Option<&SpanContext>) -> SpanData {
|
||||
let context = match parent {
|
||||
Some(p) => p.child(),
|
||||
None => SpanContext::new_root(),
|
||||
};
|
||||
|
||||
let span = SpanData {
|
||||
context: context.clone(),
|
||||
operation_name: operation_name.to_string(),
|
||||
start_time: chrono::Utc::now(),
|
||||
end_time: None,
|
||||
duration_us: None,
|
||||
status: SpanStatus::Unset,
|
||||
attributes: HashMap::new(),
|
||||
events: Vec::new(),
|
||||
};
|
||||
|
||||
// Track active span
|
||||
self.active_spans
|
||||
.write()
|
||||
.insert(context.span_id.clone(), span.clone());
|
||||
|
||||
span
|
||||
}
|
||||
|
||||
/// End a span and record it
|
||||
pub fn end_span(&self, mut span: SpanData, status: SpanStatus) {
|
||||
let end_time = chrono::Utc::now();
|
||||
span.end_time = Some(end_time);
|
||||
span.duration_us =
|
||||
Some((end_time - span.start_time).num_microseconds().unwrap_or(0) as u64);
|
||||
span.status = status;
|
||||
|
||||
// Remove from active spans
|
||||
self.active_spans.write().remove(&span.context.span_id);
|
||||
|
||||
// Add to buffer (with size limit)
|
||||
let mut buffer = self.spans_buffer.write();
|
||||
if buffer.len() >= self.config.max_spans_buffer {
|
||||
buffer.remove(0); // Remove oldest
|
||||
}
|
||||
buffer.push(span);
|
||||
}
|
||||
|
||||
/// Add an event to a span
|
||||
pub fn add_span_event(&self, span_id: &str, name: &str, attributes: HashMap<String, String>) {
|
||||
if let Some(span) = self.active_spans.write().get_mut(span_id) {
|
||||
span.events.push(SpanEvent {
|
||||
name: name.to_string(),
|
||||
timestamp: chrono::Utc::now(),
|
||||
attributes,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Set span attribute
|
||||
pub fn set_span_attribute(&self, span_id: &str, key: &str, value: &str) {
|
||||
if let Some(span) = self.active_spans.write().get_mut(span_id) {
|
||||
span.attributes.insert(key.to_string(), value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
/// Get completed spans for export
|
||||
pub fn drain_spans(&self) -> Vec<SpanData> {
|
||||
let mut buffer = self.spans_buffer.write();
|
||||
std::mem::take(&mut *buffer)
|
||||
}
|
||||
|
||||
/// Get active span count
|
||||
pub fn active_span_count(&self) -> usize {
|
||||
self.active_spans.read().len()
|
||||
}
|
||||
|
||||
/// Get buffered span count
|
||||
pub fn buffered_span_count(&self) -> usize {
|
||||
self.spans_buffer.read().len()
|
||||
}
|
||||
|
||||
/// Export spans to Jaeger format (JSON)
|
||||
pub fn export_jaeger_json(&self) -> String {
|
||||
let spans = self.spans_buffer.read();
|
||||
serde_json::to_string_pretty(&*spans).unwrap_or_else(|_| "[]".to_string())
|
||||
}
|
||||
|
||||
/// Check if telemetry is enabled
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
/// Get the service name
|
||||
pub fn service_name(&self) -> &str {
|
||||
&self.config.service_name
|
||||
}
|
||||
}
|
||||
|
||||
/// RAII guard for automatic span management
|
||||
pub struct SpanGuard<'a> {
|
||||
manager: &'a TelemetryManager,
|
||||
span: Option<SpanData>,
|
||||
}
|
||||
|
||||
impl<'a> SpanGuard<'a> {
|
||||
/// Create a new span guard
|
||||
pub fn new(
|
||||
manager: &'a TelemetryManager,
|
||||
operation_name: &str,
|
||||
parent: Option<&SpanContext>,
|
||||
) -> Self {
|
||||
let span = manager.start_span(operation_name, parent);
|
||||
Self {
|
||||
manager,
|
||||
span: Some(span),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the span context
|
||||
pub fn context(&self) -> Option<&SpanContext> {
|
||||
self.span.as_ref().map(|s| &s.context)
|
||||
}
|
||||
|
||||
/// Add an attribute to the span
|
||||
pub fn set_attribute(&self, key: &str, value: &str) {
|
||||
if let Some(span) = &self.span {
|
||||
self.manager
|
||||
.set_span_attribute(&span.context.span_id, key, value);
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark the span as successful
|
||||
pub fn set_ok(mut self) {
|
||||
if let Some(span) = self.span.take() {
|
||||
self.manager.end_span(span, SpanStatus::Ok);
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark the span as failed
|
||||
pub fn set_error(mut self, error: &str) {
|
||||
if let Some(span) = self.span.take() {
|
||||
self.manager
|
||||
.end_span(span, SpanStatus::Error(error.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SpanGuard<'_> {
|
||||
fn drop(&mut self) {
|
||||
if let Some(span) = self.span.take() {
|
||||
// If not explicitly ended, mark as unset
|
||||
self.manager.end_span(span, SpanStatus::Unset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Instrumentation macros for common operations
|
||||
#[macro_export]
|
||||
macro_rules! instrument_inference {
|
||||
($telemetry:expr, $model_name:expr, $batch_size:expr) => {{
|
||||
let guard = $crate::telemetry::SpanGuard::new($telemetry, "inference", None);
|
||||
guard.set_attribute("model", $model_name);
|
||||
guard.set_attribute("batch_size", &$batch_size.to_string());
|
||||
guard
|
||||
}};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_generate_trace_id() {
|
||||
let id1 = generate_trace_id();
|
||||
let id2 = generate_trace_id();
|
||||
assert_ne!(id1, id2);
|
||||
assert_eq!(id1.len(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_span_context_new_root() {
|
||||
let ctx = SpanContext::new_root();
|
||||
assert!(!ctx.trace_id.is_empty());
|
||||
assert!(!ctx.span_id.is_empty());
|
||||
assert!(ctx.parent_span_id.is_none());
|
||||
assert!(ctx.sampled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_span_context_child() {
|
||||
let parent = SpanContext::new_root();
|
||||
let child = parent.child();
|
||||
|
||||
assert_eq!(child.trace_id, parent.trace_id);
|
||||
assert_ne!(child.span_id, parent.span_id);
|
||||
assert_eq!(child.parent_span_id, Some(parent.span_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_traceparent_roundtrip() {
|
||||
let ctx = SpanContext::new_root();
|
||||
let header = ctx.to_traceparent();
|
||||
let parsed = SpanContext::from_traceparent(&header);
|
||||
|
||||
assert!(parsed.is_some());
|
||||
let parsed = parsed.unwrap();
|
||||
assert_eq!(parsed.trace_id, ctx.trace_id);
|
||||
assert_eq!(parsed.parent_span_id, Some(ctx.span_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trace_config_default() {
|
||||
let config = TraceConfig::default();
|
||||
assert_eq!(config.sample_rate, 1.0);
|
||||
assert!(config.enable_console);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trace_config_production() {
|
||||
let config = TraceConfig::production("test-service");
|
||||
assert_eq!(config.sample_rate, 0.1);
|
||||
assert!(!config.enable_console);
|
||||
assert!(config.enable_json);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_telemetry_manager_creation() {
|
||||
let config = TraceConfig::default();
|
||||
let manager = TelemetryManager::new(config).await.unwrap();
|
||||
assert!(manager.is_enabled());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_span_lifecycle() {
|
||||
let config = TraceConfig::default();
|
||||
let manager = TelemetryManager::new(config).await.unwrap();
|
||||
|
||||
let span = manager.start_span("test_operation", None);
|
||||
assert_eq!(manager.active_span_count(), 1);
|
||||
|
||||
manager.end_span(span, SpanStatus::Ok);
|
||||
assert_eq!(manager.active_span_count(), 0);
|
||||
assert_eq!(manager.buffered_span_count(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_span_guard() {
|
||||
let config = TraceConfig::default();
|
||||
let manager = TelemetryManager::new(config).await.unwrap();
|
||||
|
||||
{
|
||||
let guard = SpanGuard::new(&manager, "guarded_op", None);
|
||||
guard.set_attribute("key", "value");
|
||||
assert_eq!(manager.active_span_count(), 1);
|
||||
guard.set_ok();
|
||||
}
|
||||
|
||||
assert_eq!(manager.active_span_count(), 0);
|
||||
assert_eq!(manager.buffered_span_count(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
//! Tracing utilities and integration.
|
||||
Reference in New Issue
Block a user