import os def get_mpi_rank(): """ Try to get Rank ID from common environment variables. If not found, return "0". """ rank_vars = [ "OMPI_COMM_WORLD_RANK", "PMI_RANK", "SLURM_PROCID", "RANK", ] for var in rank_vars: if var in os.environ: return int(os.environ[var]) return 0