`core._rank_utils`#

Low-level rank utilities with minimal dependencies to avoid circular imports.

Module Contents#

Functions#

`safe_get_rank`	Get the distributed rank safely, even if torch.distributed is not initialized.
`safe_get_world_size`	Get the distributed world size safely, even if torch.distributed is not initialized.
`set_default_log_ranks`	Set the ranks that `log_single_rank` and `warn_single_rank` write on by default.
`get_default_log_ranks`	Return the ranks that the single-rank logging helpers write on by default.
`log_single_rank`	Log a message only on a single rank.
`warn_single_rank`	Issue a warning only on a single rank.

Data#

_DEFAULT_LOG_RANKS

API#

core._rank_utils.safe_get_rank() → int#

Get the distributed rank safely, even if torch.distributed is not initialized.

Fallback order:

torch.distributed.get_rank() (if initialized)
RANK environment variable (torchrun/torchelastic)
SLURM_PROCID environment variable (SLURM)
Default: 0 (with warning)

Returns:: The rank of the current process.
Return type:: int

core._rank_utils.safe_get_world_size() → int#

Get the distributed world size safely, even if torch.distributed is not initialized.

Fallback order:

torch.distributed.get_world_size() (if initialized)
WORLD_SIZE environment variable (torchrun/torchelastic)
SLURM_NTASKS environment variable (SLURM)
Default: 1 (with warning)

Returns:: The total number of processes in the distributed job.

core._rank_utils._DEFAULT_LOG_RANKS: tuple[int, ...]#: (0,)

core._rank_utils.set_default_log_ranks(ranks: collections.abc.Iterable[int]) → None#

Set the ranks that log_single_rank and warn_single_rank write on by default.

Call once, after torch distributed is initialized and before the model is built, so that setup-time messages are covered. A call site that passes rank explicitly is unaffected.

Parameters:: ranks – Ranks to log on. Duplicates are ignored.

core._rank_utils.get_default_log_ranks() → tuple[int, ...]#: Return the ranks that the single-rank logging helpers write on by default.

core._rank_utils.log_single_rank(

logger: logging.Logger,

level: int,

msg: object,

*args: Any,

rank: Optional[int] = None,

**kwargs: Any,

) → None#

Log a message only on a single rank.

If torch distributed is initialized, write log on only one rank.

Parameters:

logger – The logger to write the logs.
level – Logging level for the message.
msg – Message format string.
*args – Message format arguments.
rank – The rank to write on. Defaults to None, meaning the ranks configured by set_default_log_ranks (rank 0 unless it has been changed).
**kwargs – Additional logging.Logger.log keyword arguments.

core._rank_utils.warn_single_rank( message: str, category: type[Warning] = UserWarning, stacklevel: int = 2, rank: Optional[int] = None, ) → None#

Issue a warning only on a single rank.

Use for warnings that describe a property of the job rather than of the calling rank, such as deprecated settings and experimental-API notices. Every rank raises those identically, so a large job repeats one message thousands of times in a shared log.

safe_get_rank reads the RANK or SLURM_PROCID environment variable when torch distributed is not initialized, so this also works at import time.