utils.py

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