espnet2.train.distributed_utils.get_rank

source

espnet2.train.distributed_utils.get_rank(prior=None, launcher: str | None = None) → int | None

Get the rank of the current process in a distributed setting.

This function retrieves the rank of the current process. The rank is defined as an integer that uniquely identifies a process within a distributed training setup. If the prior argument is provided, it will be used as the rank. If prior is not specified, the function checks the environment variables or uses the specified launcher (e.g., “slurm”, “mpi”) to determine the rank.

• Parameters:
  • prior (Optional *[*int ]) – An optional prior rank to use if provided.
  • launcher (Optional *[*str ]) – The launcher used to start the process (e.g., “slurm”, “mpi”).
• Returns: The rank of the current process, or None if it cannot be determined.
• Return type: Optional[int]
• Raises:
  • RuntimeError – If the process is launched with an unsupported
  • launcher or if the required environment variables are not set. –

Examples

>>> os.environ["RANK"] = "2"
>>> get_rank()  # Returns 2

>>> os.environ["SLURM_PROCID"] = "0"
>>> get_rank(launcher="slurm")  # Returns 0

>>> get_rank(1)  # Returns 1

NOTE

If prior is None and no environment variable is set, the function will return None.