burn_telemetry 0.1.0

//! Multi-logger wrapper for combining multiple MetricLogger implementations

use burn::train::{
    logger::MetricLogger,
    metric::{MetricEntry, NumericEntry},
};

/// A wrapper that allows multiple loggers to be used simultaneously
///
/// This enables logging to multiple backends at once, such as both
/// OpenTelemetry and Rerun during the same training run.
///
/// # Example
///
/// ```no_run
/// # #[cfg(all(feature = "otel", feature = "rerun"))]
/// # {
/// use burn_telemetry::MultiLogger;
/// use burn_telemetry::otel::{OpenTelemetryConfig, OpenTelemetryLogger};
/// use burn_telemetry::rerun::{RerunConfig, RerunLogger};
///
/// let otel_logger = OpenTelemetryLogger::new(
///     OpenTelemetryConfig::default(),
///     "job",
///     Some("device-0"),
/// ).expect("Failed to create OTel logger");
///
/// let rerun_logger = RerunLogger::new(
///     RerunConfig::default(),
///     "job",
/// ).expect("Failed to create Rerun logger");
///
/// let multi_logger = MultiLogger::new()
///     .add_logger(otel_logger)
///     .add_logger(rerun_logger);
/// # }
/// ```
pub struct MultiLogger {
    loggers: Vec<Box<dyn MetricLogger>>,
}

impl MultiLogger {
    /// Create a new MultiLogger with no loggers
    pub fn new() -> Self {
        Self {
            loggers: Vec::new(),
        }
    }

    /// Add a logger to the multi-logger
    pub fn add_logger<L: MetricLogger + 'static>(mut self, logger: L) -> Self {
        self.loggers.push(Box::new(logger));
        self
    }

    /// Add a boxed logger to the multi-logger
    pub fn add_boxed_logger(mut self, logger: Box<dyn MetricLogger>) -> Self {
        self.loggers.push(logger);
        self
    }
}

impl Default for MultiLogger {
    fn default() -> Self {
        Self::new()
    }
}

impl MetricLogger for MultiLogger {
    fn log(&mut self, item: &MetricEntry) {
        for logger in &mut self.loggers {
            logger.log(item);
        }
    }

    fn end_epoch(&mut self, epoch: usize) {
        for logger in &mut self.loggers {
            logger.end_epoch(epoch);
        }
    }

    fn read_numeric(&mut self, name: &str, epoch: usize) -> Result<Vec<NumericEntry>, String> {
        // Return the result from the first logger that supports reading
        for logger in &mut self.loggers {
            if let Ok(result) = logger.read_numeric(name, epoch) {
                return Ok(result);
            }
        }
        Err("No logger in MultiLogger supports reading metrics".to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    };

    // Mock logger for testing
    struct MockLogger {
        log_count: Arc<AtomicUsize>,
        end_epoch_count: Arc<AtomicUsize>,
    }

    impl MockLogger {
        fn new() -> Self {
            Self::with_counters(
                Arc::new(AtomicUsize::new(0)),
                Arc::new(AtomicUsize::new(0)),
            )
        }

        fn with_counters(
            log_count: Arc<AtomicUsize>,
            end_epoch_count: Arc<AtomicUsize>,
        ) -> Self {
            Self {
                log_count,
                end_epoch_count,
            }
        }
    }

    impl MetricLogger for MockLogger {
        fn log(&mut self, _item: &MetricEntry) {
            self.log_count.fetch_add(1, Ordering::SeqCst);
        }

        fn end_epoch(&mut self, _epoch: usize) {
            self.end_epoch_count.fetch_add(1, Ordering::SeqCst);
        }

        fn read_numeric(
            &mut self,
            _name: &str,
            _epoch: usize,
        ) -> Result<Vec<NumericEntry>, String> {
            Err("Mock logger does not support reading".to_string())
        }
    }

    #[test]
    fn test_multi_logger_creation() {
        let logger = MultiLogger::new();
        assert_eq!(logger.loggers.len(), 0);
    }

    #[test]
    fn test_multi_logger_add_loggers() {
        let logger = MultiLogger::new()
            .add_logger(MockLogger::new())
            .add_logger(MockLogger::new());
        assert_eq!(logger.loggers.len(), 2);
    }

    #[test]
    fn test_multi_logger_logs_to_all() {
        let first_log = Arc::new(AtomicUsize::new(0));
        let first_end = Arc::new(AtomicUsize::new(0));
        let second_log = Arc::new(AtomicUsize::new(0));
        let second_end = Arc::new(AtomicUsize::new(0));

        let mut logger = MultiLogger::new()
            .add_logger(MockLogger::with_counters(first_log.clone(), first_end.clone()))
            .add_logger(MockLogger::with_counters(second_log.clone(), second_end.clone()));

        let entry = MetricEntry::new(
            Arc::new("test".to_string()),
            "1.0".to_string(),
            "1.0".to_string(),
        );

        logger.log(&entry);
        logger.end_epoch(1);

        assert_eq!(first_log.load(Ordering::SeqCst), 1);
        assert_eq!(first_end.load(Ordering::SeqCst), 1);
        assert_eq!(second_log.load(Ordering::SeqCst), 1);
        assert_eq!(second_end.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_multi_logger_read_not_supported() {
        let mut logger = MultiLogger::new().add_logger(MockLogger::new());

        let result = logger.read_numeric("test", 1);
        assert!(result.is_err());
    }
}