burn_telemetry 0.1.0

//! OpenTelemetry metric logger implementation

use super::OpenTelemetryConfig;
use crate::metric::{metric_identifier, numeric_value};
use burn::train::{
    logger::MetricLogger,
    metric::{MetricEntry, NumericEntry},
};
use opentelemetry::{
    metrics::{Gauge, Meter, MeterProvider},
    KeyValue,
};
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{
    metrics::{PeriodicReader, SdkMeterProvider},
    Resource,
};
use std::{collections::HashMap, thread};
use tokio::runtime::Handle;

/// OpenTelemetry metric logger for Burn training
///
/// This logger exports training metrics to OpenTelemetry-compatible backends.
/// It supports both OTLP and stdout exporters and can be configured for
/// minimal overhead and multi-GPU/node scenarios.
pub struct OpenTelemetryLogger {
    meter: Meter,
    gauges: HashMap<String, Gauge<f64>>,
    _provider: SdkMeterProvider,
    metric_dimensions: Vec<KeyValue>,
    current_epoch: usize,
    flush_threshold: usize,
    pending_samples: usize,
    async_flush: bool,
}

impl OpenTelemetryLogger {
    /// Create a new OpenTelemetry logger
    ///
    /// # Arguments
    ///
    /// * `config` - Configuration for the logger
    /// * `job_name` - Name of the training job
    /// * `device_id` - Optional device identifier for multi-GPU setups
    ///
    /// # Returns
    ///
    /// A new logger instance or an error if initialization fails
    pub fn new(
        config: OpenTelemetryConfig,
        job_name: impl Into<String>,
        device_id: Option<impl Into<String>>,
    ) -> Result<Self, String> {
        let job_name = job_name.into();
        let device_id = device_id.map(|d| d.into());
        let OpenTelemetryConfig {
            endpoint,
            export_interval,
            export_timeout,
            service_name,
            resource_attributes,
            max_batch_size,
            async_export,
            ..
        } = config;

        // Build resource with service name and custom attributes
        let mut resource_kvs = vec![
            KeyValue::new("service.name", service_name),
            KeyValue::new("job.name", job_name.clone()),
        ];

        if let Some(ref dev_id) = device_id {
            resource_kvs.push(KeyValue::new("device.id", dev_id.clone()));
        }

        for (key, value) in resource_attributes {
            resource_kvs.push(KeyValue::new(key, value));
        }

        let mut metric_dimensions = vec![KeyValue::new("job", job_name.clone())];
        if let Some(ref dev_id) = device_id {
            metric_dimensions.push(KeyValue::new("device", dev_id.clone()));
        }

        let resource = Resource::builder_empty()
            .with_attributes(resource_kvs)
            .build();

        // Create reader with appropriate exporter
        let mut provider_builder = SdkMeterProvider::builder().with_resource(resource);

        provider_builder = if let Some(endpoint) = endpoint {
            // Use OTLP exporter
            let exporter = opentelemetry_otlp::MetricExporter::builder()
                .with_tonic()
                .with_endpoint(endpoint)
                .with_timeout(export_timeout)
                .build()
                .map_err(|e| format!("Failed to create OTLP exporter: {}", e))?;

            provider_builder.with_reader(
                PeriodicReader::builder(exporter)
                    .with_interval(export_interval)
                    .build(),
            )
        } else {
            // Use stdout exporter for testing/debugging
            let exporter = opentelemetry_stdout::MetricExporter::default();
            provider_builder.with_reader(
                PeriodicReader::builder(exporter)
                    .with_interval(export_interval)
                    .build(),
            )
        };

        // Create meter provider
        let provider = provider_builder.build();

        let meter = provider.meter("burn_telemetry");

        Ok(Self {
            meter,
            gauges: HashMap::new(),
            _provider: provider,
            metric_dimensions,
            current_epoch: 1,
            flush_threshold: max_batch_size,
            pending_samples: 0,
            async_flush: async_export,
        })
    }

    /// Get or create a gauge for a metric
    fn get_or_create_gauge(&mut self, name: &str) -> Gauge<f64> {
        if let Some(gauge) = self.gauges.get(name) {
            return gauge.clone();
        }

        let gauge = self
            .meter
            .f64_gauge(metric_identifier(name))
            .with_description(format!("Burn training metric: {}", name))
            .build();
        self.gauges.insert(name.to_string(), gauge.clone());
        gauge
    }

    /// Get or create a counter for a metric
    #[allow(dead_code)]
    fn get_or_create_counter(&self, name: &str) -> opentelemetry::metrics::Counter<u64> {
        self.meter
            .u64_counter(metric_identifier(name))
            .with_description(format!("Burn training metric: {}", name))
            .build()
    }
    /// Build attributes for a metric
    fn build_attributes(&self) -> Vec<KeyValue> {
        let mut attributes = Vec::with_capacity(self.metric_dimensions.len() + 1);
        attributes.push(KeyValue::new("epoch", self.current_epoch as i64));
        attributes.extend(self.metric_dimensions.iter().cloned());
        attributes
    }

    fn flush_if_needed(&mut self) {
        if self.flush_threshold == 0 {
            return;
        }
        if self.pending_samples < self.flush_threshold {
            return;
        }
        self.pending_samples = 0;
        if self.async_flush {
            self.flush_async();
        } else {
            Self::flush_provider(self._provider.clone());
        }
    }

    fn flush_async(&self) {
        let provider = self._provider.clone();
        if let Ok(handle) = Handle::try_current() {
            handle.spawn_blocking(move || Self::flush_provider(provider));
        } else {
            thread::spawn(move || Self::flush_provider(provider));
        }
    }

    fn flush_provider(provider: SdkMeterProvider) {
        if let Err(err) = provider.force_flush() {
            eprintln!("OpenTelemetry flush failed: {err}");
        }
    }
}

impl MetricLogger for OpenTelemetryLogger {
    fn log(&mut self, item: &MetricEntry) {
        if let Some(value) = numeric_value(item) {
            let gauge = self.get_or_create_gauge(&item.name);
            let attributes = self.build_attributes();
            gauge.record(value, &attributes);
            self.pending_samples += 1;
            self.flush_if_needed();
        }
    }

    fn end_epoch(&mut self, epoch: usize) {
        self.current_epoch = epoch + 1;
    }

    fn read_numeric(&mut self, _name: &str, _epoch: usize) -> Result<Vec<NumericEntry>, String> {
        // OpenTelemetry is write-only; we don't support reading back
        Err("OpenTelemetry logger does not support reading metrics".to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn::train::metric::MetricEntry;
    use std::sync::Arc;

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

    #[test]
    fn test_metric_attributes_include_epoch_and_device() {
        let config = OpenTelemetryConfig::default();
        let mut logger =
            OpenTelemetryLogger::new(config, "test-job", Some("gpu-0")).expect("logger");

        logger.current_epoch = 3;
        let attributes = logger.build_attributes();

        assert!(
            attributes.contains(&KeyValue::new("epoch", 3_i64)),
            "epoch attribute missing: {attributes:?}"
        );
        assert!(
            attributes.contains(&KeyValue::new("device", "gpu-0")),
            "device attribute missing: {attributes:?}"
        );
        assert!(
            attributes.contains(&KeyValue::new("job", "test-job")),
            "job attribute missing: {attributes:?}"
        );
    }

    #[test]
    fn test_gauge_cache_reuse() {
        let config = OpenTelemetryConfig::default();
        let mut logger =
            OpenTelemetryLogger::new(config, "job", Option::<String>::None).expect("logger");

        logger.get_or_create_gauge("My Metric");
        logger.get_or_create_gauge("My Metric");

        assert_eq!(logger.gauges.len(), 1);
    }

    #[test]
    fn test_batching_resets_pending_counter() {
        let config = OpenTelemetryConfig::new("job")
            .with_max_batch_size(2)
            .with_async_export(false);
        let mut logger =
            OpenTelemetryLogger::new(config, "job", Option::<String>::None).expect("logger");

        let entry = MetricEntry::new(
            Arc::new("loss".to_string()),
            "0.1".to_string(),
            "0.1".to_string(),
        );

        assert_eq!(logger.pending_samples, 0);
        logger.log(&entry);
        assert_eq!(logger.pending_samples, 1);
        logger.log(&entry);
        assert_eq!(logger.pending_samples, 0);
    }
}