utils/distributed_utils.py

import os

# Check standard environment variables for distributed training
# Default to True (rank 0) if not in a distributed environment
_rank = int(os.environ.get("RANK", 0))
_local_rank = int(os.environ.get("LOCAL_RANK", 0))

# We consider it rank zero if global rank is 0.
# Local rank check is usually redundant if rank is 0, but good for sanity.
is_rank_zero = _rank == 0