setup_logging.py

import logging
import sys
import traceback
import warnings
import os


class ContextFilter(logging.Filter):
    def filter(self, record):
        record.notebook = NOTEBOOK_NAME
        record.directory = CURRENT_WORKING_DIR
        return True


# Define a custom log format
formatter = logging.Formatter(
    '%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - '
    '[Notebook: %(notebook)s] - [Directory: %(directory)s] - %(message)s'
)

# Retrieve notebook name and current working directory from environment variables
NOTEBOOK_NAME = os.getenv('CAL_NOTEBOOK_NAME', 'Unknown notebook')
CURRENT_WORKING_DIR = os.getenv('CAL_WORKING_DIR', 'Unknown directory')

# Create a logger
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)  # Capture all levels

# Create file handlers
error_handler = logging.FileHandler('errors.log')
error_handler.setLevel(logging.ERROR)
error_handler.setFormatter(formatter)

warning_handler = logging.FileHandler('warnings.log')
warning_handler.setLevel(logging.WARNING)
# Avoid errors being logged in warnings.log
warning_handler.addFilter(lambda record: record.levelno < logging.ERROR)
warning_handler.setFormatter(formatter)

# Create console handler for info and debug
# No formatter used for console handler on purpose.
# console is slurm.out & the notebook output cells.
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.DEBUG)

# Add the custom filter to handlers
context_filter = ContextFilter()
error_handler.addFilter(context_filter)
warning_handler.addFilter(context_filter)
console_handler.addFilter(context_filter)

# Add handlers to logger
logger.addHandler(error_handler)
logger.addHandler(warning_handler)
logger.addHandler(console_handler)

handling_error = False


def safe_handle_error(exc_type, exc_value, exc_traceback):
    global handling_error
    if handling_error:  # Avoid infinite loop of errors.
        sys.stderr.write("Recursive error detected!\n")
        traceback.print_exception(exc_type, exc_value, exc_traceback, file=sys.stderr)
        return

    handling_error = True
    try:
        # Log the error with the notebook name and additional metadata
        logger.error(
            "An error occurred. Exception type: %s, Message: %s",
            exc_type.__name__, str(exc_value),
            extra={'notebook': NOTEBOOK_NAME, 'directory': CURRENT_WORKING_DIR},
            exc_info=(exc_type, exc_value, exc_traceback)
        )

    except Exception as log_error:
        sys.stderr.write(f"Logging failed: {log_error}\n")
        traceback.print_exception(exc_type, exc_value, exc_traceback, file=sys.stderr)
    finally:
        handling_error = False


def handle_warning(message, category, filename, lineno, file=None, line=None):
    try:
        # Log the warning with the notebook name and additional metadata
        logger.warning(
            "Warning occurred: %s, File: %s, Line: %d",
            message, filename, lineno,
            extra={'notebook': NOTEBOOK_NAME, 'directory': CURRENT_WORKING_DIR}
        )
    except Exception as log_error:
        sys.stderr.write(f"Logging failed: {log_error}\n")


# Replace the handlers with our custom ones
sys.excepthook = safe_handle_error
warnings.showwarning = handle_warning

# Override IPython's exception handling
import IPython
IPython.core.interactiveshell.InteractiveShell.showtraceback = lambda self, *args, **kwargs: safe_handle_error(*sys.exc_info())