burn_telemetry 0.1.0

//! Configuration for OpenTelemetry logger

use std::time::Duration;

/// Configuration for OpenTelemetry metric export
#[derive(Debug, Clone)]
pub struct OpenTelemetryConfig {
    /// OTLP endpoint URL (e.g., "http://localhost:4317")
    pub endpoint: Option<String>,

    /// Export interval for batched metrics
    pub export_interval: Duration,

    /// Maximum number of metric samples to buffer before forcing a flush.
    ///
    /// Each time the logger records a numeric value, it increments a lightweight counter.
    /// When the counter reaches this threshold, the underlying meter provider is flushed.
    pub max_batch_size: usize,

    /// Whether flushes should run asynchronously.
    ///
    /// When enabled, flushes are offloaded to the current Tokio runtime (or a background thread
    /// if no runtime is available) so the training loop never blocks on network I/O.
    pub async_export: bool,

    /// Timeout for export operations
    pub export_timeout: Duration,

    /// Service name for resource attribution
    pub service_name: String,

    /// Additional resource attributes (e.g., for multi-GPU/node tracking)
    pub resource_attributes: Vec<(String, String)>,
}

impl Default for OpenTelemetryConfig {
    fn default() -> Self {
        Self {
            endpoint: None, // Use stdout exporter by default
            export_interval: Duration::from_secs(10),
            max_batch_size: 100,
            async_export: true,
            export_timeout: Duration::from_secs(5),
            service_name: "burn-training".to_string(),
            resource_attributes: Vec::new(),
        }
    }
}

impl OpenTelemetryConfig {
    /// Create a new configuration with the given service name
    pub fn new(service_name: impl Into<String>) -> Self {
        Self {
            service_name: service_name.into(),
            ..Default::default()
        }
    }

    /// Set the OTLP endpoint
    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.endpoint = Some(endpoint.into());
        self
    }

    /// Set the export interval
    pub fn with_export_interval(mut self, interval: Duration) -> Self {
        self.export_interval = interval;
        self
    }

    /// Set the maximum batch size.
    ///
    /// A value of `0` disables forced flushes and leaves export scheduling entirely
    /// to the OpenTelemetry SDK.
    pub fn with_max_batch_size(mut self, size: usize) -> Self {
        self.max_batch_size = size;
        self
    }

    /// Add a resource attribute for multi-GPU/node tracking
    pub fn with_resource_attribute(
        mut self,
        key: impl Into<String>,
        value: impl Into<String>,
    ) -> Self {
        self.resource_attributes.push((key.into(), value.into()));
        self
    }

    /// Set whether to use async export.
    ///
    /// When disabled, flushes are performed synchronously on the training thread.
    pub fn with_async_export(mut self, async_export: bool) -> Self {
        self.async_export = async_export;
        self
    }

    /// Set the export timeout
    pub fn with_export_timeout(mut self, timeout: Duration) -> Self {
        self.export_timeout = timeout;
        self
    }
}

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

    #[test]
    fn test_default_config() {
        let config = OpenTelemetryConfig::default();
        assert_eq!(config.service_name, "burn-training");
        assert!(config.async_export);
        assert_eq!(config.max_batch_size, 100);
    }

    #[test]
    fn test_builder_pattern() {
        let config = OpenTelemetryConfig::new("test-service")
            .with_endpoint("http://localhost:4317")
            .with_max_batch_size(50)
            .with_resource_attribute("gpu.id", "0")
            .with_async_export(false);

        assert_eq!(config.service_name, "test-service");
        assert_eq!(config.endpoint, Some("http://localhost:4317".to_string()));
        assert_eq!(config.max_batch_size, 50);
        assert!(!config.async_export);
        assert_eq!(config.resource_attributes.len(), 1);
    }
}