| 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) | |