utils/logging_utils.py

import inspect
import logging
import os


def _process_index() -> int:
    """Return torch.distributed rank, falling back to env vars or 0."""
    try:
        import torch.distributed as dist
        if dist.is_available() and dist.is_initialized():
            return dist.get_rank()
    except Exception:
        pass
    return int(os.environ.get("RANK", "0"))


def log_for_0(msg, *args, level=logging.INFO):
    """Log only on the first process (rank == 0)."""
    if _process_index() != 0:
        return
    caller_module = inspect.currentframe().f_back.f_globals.get("__name__", __name__)
    logging.getLogger(caller_module).log(level, msg, *args)