File size: 658 Bytes
6b56858 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | 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)
|