examples/enhanced_logging_example.py

#!/usr/bin/env python3
"""
Example script demonstrating enhanced logging features in F5-TTS trainer.

This script shows how to:
1. Use different logging configurations
2. Monitor training progress with detailed logs
3. Handle errors gracefully with proper logging
4. Use predefined logging configurations
"""

import os
import sys
import time
from datetime import datetime

# Add the src directory to the path so we can import f5_tts
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from f5_tts.model.logging_config import (
    F5TTSLoggingConfig, 
    setup_advanced_logger, 
    get_predefined_config,
    log_training_start,
    log_training_end,
    log_performance_metrics
)


def example_basic_logging():
    """Example of basic logging setup."""
    print("=== Basic Logging Example ===")
    
    # Create a simple logging configuration
    config = F5TTSLoggingConfig(
        log_level="INFO",
        console_level="INFO",
        file_level="DEBUG",
        enable_console=True,
        enable_file=True,
        enable_error_file=True
    )
    
    # Setup logger
    checkpoint_path = "examples/logs/basic_example"
    logger = setup_advanced_logger(checkpoint_path, config)
    
    # Simulate training start
    training_config = {
        "start_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "checkpoint_path": checkpoint_path,
        "device": "cuda:0",
        "mixed_precision": "fp16"
    }
    
    log_training_start(logger, training_config)
    
    # Simulate some training steps
    for step in range(5):
        logger.info(f"Training step {step + 1}")
        logger.debug(f"Detailed debug info for step {step + 1}")
        
        if step == 2:
            logger.warning("This is a warning message")
        
        time.sleep(0.1)  # Simulate processing time
    
    # Simulate training end
    training_stats = {
        "total_time": "00:01:30",
        "final_loss": 0.123456,
        "best_loss": 0.098765,
        "total_updates": 1000,
        "total_batches": 5000
    }
    
    log_training_end(logger, training_stats)
    print("Basic logging example completed. Check the logs directory.")


def example_verbose_logging():
    """Example of verbose logging for debugging."""
    print("\n=== Verbose Logging Example ===")
    
    # Use predefined verbose configuration
    config = get_predefined_config("verbose")
    
    checkpoint_path = "examples/logs/verbose_example"
    logger = setup_advanced_logger(checkpoint_path, config)
    
    logger.info("Starting verbose logging example")
    
    # Simulate detailed training monitoring
    for epoch in range(3):
        logger.info(f"Starting epoch {epoch + 1}")
        
        for batch in range(10):
            # Simulate performance metrics
            metrics = {
                "current_loss": 0.1 + (batch * 0.01),
                "avg_loss": 0.15 + (batch * 0.005),
                "learning_rate": 1e-4 * (0.9 ** epoch),
                "batch_time": 0.05 + (batch * 0.001),
                "memory_usage": f"{1024 + batch * 50}MB"
            }
            
            if batch % 3 == 0:
                log_performance_metrics(logger, metrics)
            
            logger.debug(f"Batch {batch} processed successfully")
        
        logger.info(f"Epoch {epoch + 1} completed")
    
    logger.info("Verbose logging example completed")


def example_production_logging():
    """Example of production-ready logging configuration."""
    print("\n=== Production Logging Example ===")
    
    # Use production configuration
    config = get_predefined_config("production")
    
    checkpoint_path = "examples/logs/production_example"
    logger = setup_advanced_logger(checkpoint_path, config)
    
    logger.info("Starting production logging example")
    
    # Simulate production training with error handling
    try:
        for step in range(10):
            if step == 5:
                # Simulate an error
                raise RuntimeError("Simulated training error")
            
            logger.info(f"Processing step {step + 1}")
            time.sleep(0.05)
            
    except Exception as e:
        logger.error(f"Training failed: {str(e)}")
        logger.error("Attempting to save checkpoint before exit")
    
    logger.info("Production logging example completed")


def example_custom_logging():
    """Example of custom logging configuration."""
    print("\n=== Custom Logging Example ===")
    
    # Create custom configuration
    custom_config = F5TTSLoggingConfig(
        log_level="DEBUG",
        console_level="INFO",
        file_level="DEBUG",
        log_format="custom",
        custom_formatters={
            "default": "%(asctime)s | %(levelname)-8s | %(message)s",
            "error": "%(asctime)s | ERROR | %(name)s:%(lineno)d | %(message)s"
        },
        max_file_size_mb=10,
        backup_count=3
    )
    
    checkpoint_path = "examples/logs/custom_example"
    logger = setup_advanced_logger(checkpoint_path, custom_config)
    
    logger.info("Custom logging configuration active")
    logger.debug("This is a debug message with custom formatting")
    logger.warning("This is a warning with custom formatting")
    logger.error("This is an error with custom formatting")
    
    logger.info("Custom logging example completed")


def example_error_handling():
    """Example of error handling with logging."""
    print("\n=== Error Handling Example ===")
    
    config = F5TTSLoggingConfig(
        log_level="INFO",
        console_level="WARNING",  # Only show warnings and errors in console
        file_level="DEBUG",       # But log everything to file
        enable_console=True,
        enable_file=True,
        enable_error_file=True
    )
    
    checkpoint_path = "examples/logs/error_example"
    logger = setup_advanced_logger(checkpoint_path, config)
    
    logger.info("Starting error handling example")
    
    # Simulate different types of errors
    try:
        # Simulate a data loading error
        logger.warning("Data loading taking longer than expected")
        time.sleep(0.1)
        
        # Simulate a model error
        raise ValueError("Invalid model configuration")
        
    except ValueError as e:
        logger.error(f"Configuration error: {str(e)}")
        logger.info("Attempting to use default configuration")
        
    try:
        # Simulate a training error
        raise RuntimeError("GPU out of memory")
        
    except RuntimeError as e:
        logger.error(f"Training error: {str(e)}")
        logger.info("Attempting to reduce batch size")
        
    logger.info("Error handling example completed")


def main():
    """Run all logging examples."""
    print("F5-TTS Enhanced Logging Examples")
    print("=" * 50)
    
    # Create logs directory
    os.makedirs("examples/logs", exist_ok=True)
    
    # Run examples
    example_basic_logging()
    example_verbose_logging()
    example_production_logging()
    example_custom_logging()
    example_error_handling()
    
    print("\n" + "=" * 50)
    print("All examples completed!")
    print("Check the 'examples/logs' directory for generated log files.")
    print("\nLog files created:")
    print("- basic_example/training.log")
    print("- basic_example/errors.log")
    print("- verbose_example/training.log")
    print("- production_example/training.log")
    print("- custom_example/training.log")
    print("- error_example/training.log")
    print("- error_example/errors.log")


if __name__ == "__main__":
    main()