src/ahdcma/utils/logging.py

"""Rich-formatted project logger with run-id binding.

Every entry-point script should call :func:`setup_logging` exactly once,
passing the run identifier produced by :func:`make_run_id`. The returned
logger writes to stderr via :class:`rich.logging.RichHandler` and to the
per-run log file under ``outputs/logs/{run_id}.log``.
"""

from __future__ import annotations

import logging
from datetime import UTC, datetime
from pathlib import Path

from rich.logging import RichHandler


def make_run_id(algo: str, task: str, seed: int, *, dim: int | None = None) -> str:
    """Return a canonical run identifier.

    Format: ``{algo}_{task}[_d{dim}]_seed{N}_{YYYYMMDD-HHMMSS}``.

    Parameters
    ----------
    algo : str
        Algorithm name (e.g., ``"ahdcma"``).
    task : str
        Task or benchmark name (e.g., ``"cifar100_vit"``).
    seed : int
        Run seed.
    dim : int, optional
        Dimensionality. When set, included in the run id so the same
        ``(algo, task, seed)`` tuple at different dimensions produces
        distinct identifiers. Required for the CEC-2022 sweep where
        each function is run at multiple dimensions.
    """
    timestamp = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
    dim_part = f"_d{dim}" if dim is not None else ""
    return f"{algo}_{task}{dim_part}_seed{seed}_{timestamp}"


class _RunIdFilter(logging.Filter):
    """Inject ``run_id`` into every record so the format string can use it."""

    def __init__(self, run_id: str) -> None:
        super().__init__()
        self._run_id = run_id

    def filter(self, record: logging.LogRecord) -> bool:
        record.run_id = self._run_id
        return True


def setup_logging(
    run_id: str,
    log_dir: str | Path = "outputs/logs",
    *,
    level: int = logging.INFO,
    verbose: bool = False,
) -> logging.Logger:
    """Configure the root ``ahdcma`` logger and return it.

    Parameters
    ----------
    run_id : str
        Run identifier; the file handler writes to ``{log_dir}/{run_id}.log``.
    log_dir : str or Path
        Directory for the per-run log file. Created if missing.
    level : int
        Default logging level for the console handler.
    verbose : bool
        If ``True``, drop the console level to ``DEBUG``.

    Returns
    -------
    logging.Logger
        The configured ``ahdcma`` logger. Calling this twice with the same
        ``run_id`` is idempotent.
    """
    log_path = Path(log_dir)
    log_path.mkdir(parents=True, exist_ok=True)

    logger = logging.getLogger("ahdcma")
    logger.setLevel(logging.DEBUG)
    logger.handlers.clear()

    fmt = "[%(asctime)s] [%(levelname)s] [run_id=%(run_id)s] %(message)s"
    datefmt = "%Y-%m-%d %H:%M:%S"

    console = RichHandler(rich_tracebacks=True, show_time=False, show_path=False)
    console.setLevel(logging.DEBUG if verbose else level)
    console.setFormatter(logging.Formatter(fmt, datefmt=datefmt))

    file_handler = logging.FileHandler(log_path / f"{run_id}.log")
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(logging.Formatter(fmt, datefmt=datefmt))

    run_filter = _RunIdFilter(run_id)
    for handler in (console, file_handler):
        handler.addFilter(run_filter)
        logger.addHandler(handler)

    logger.propagate = False
    return logger