speechbrain.utils.distributed module

Guard for running certain operations on main process only

Authors:

• Abdel Heba 2020

• Aku Rouhe 2020

• Peter Plantinga 2023

• Adel Moumen 2024

Summary

Classes:

MainProcessContext

Context manager to ensure code runs only on the main process.

Functions:

ddp_barrier	Synchronize all processes in distributed data parallel (DDP) mode.
ddp_broadcast	In DDP mode, this function will broadcast an object to all processes.
ddp_init_group	This function will initialize the ddp group if distributed_launch bool is given in the python command line.
get_rank	Get the rank of the current process.
if_main_process	Returns whether the current process is the main process.
is_distributed_initialized	Returns whether the current system is distributed.
main_process_only	Function decorator to ensure the function runs only on the main process.
rank_prefixed_message	Prefix a message with the rank of the process.
run_on_main	Runs a function with DPP (multi-gpu) support.

Reference

speechbrain.utils.distributed.rank_prefixed_message(message: str) → str

Prefix a message with the rank of the process.

Parameters:

message (str) – The message to prefix.

Returns:

The message prefixed with the rank, if known.

Return type:

str

speechbrain.utils.distributed.get_rank() → int | None

Get the rank of the current process.

This code is taken from the Pytorch Lightning library: https://github.com/Lightning-AI/pytorch-lightning/blob/bc3c9c536dc88bfa9a46f63fbce22b382a86a9cb/src/lightning/fabric/utilities/rank_zero.py#L39-L48

Returns:

The rank of the current process, or None if the rank could not be determined.

Return type:

int or None

speechbrain.utils.distributed.run_on_main(func, args=None, kwargs=None, post_func=None, post_args=None, post_kwargs=None, run_post_on_main=False)

Runs a function with DPP (multi-gpu) support.

The main function is only run on the main process. A post_function can be specified, to be on non-main processes after the main func completes. This way whatever the main func produces can be loaded on the other processes.

Parameters:

• func (callable) – Function to run on the main process.

• args (list, None) – Positional args to pass to func.

• kwargs (dict, None) – Keyword args to pass to func.

• post_func (callable, None) – Function to run after func has finished on main. By default only run on non-main processes.

• post_args (list, None) – Positional args to pass to post_func.

• post_kwargs (dict, None) – Keyword args to pass to post_func.

• run_post_on_main (bool) – Whether to run post_func on main process as well. (default: False)

speechbrain.utils.distributed.is_distributed_initialized() → bool

Returns whether the current system is distributed.

speechbrain.utils.distributed.if_main_process() → bool

Returns whether the current process is the main process.

class speechbrain.utils.distributed.MainProcessContext

Bases: object

Context manager to ensure code runs only on the main process. This is useful to make sure that MAIN_PROC_ONLY global variable is decreased even if there’s an exception raised inside of main_proc_wrapped_func fn.

__enter__()

Enter the context. Increase the counter.

__exit__(exc_type, exc_value, traceback)

Exit the context. Decrease the counter.

speechbrain.utils.distributed.main_process_only(function)

Function decorator to ensure the function runs only on the main process. This is useful for things like saving to the filesystem or logging to a web address where you only want it to happen on a single process.

speechbrain.utils.distributed.ddp_barrier()

Synchronize all processes in distributed data parallel (DDP) mode.

This function blocks the execution of the current process until all processes in the distributed group have reached the same point. It ensures that no process moves ahead until every other process has also reached this barrier. If DDP is not being used (i.e., only one process is running), this function has no effect and immediately returns.

Return type:

None

Example

>>> ddp_barrier()
>>> print("hello world")
hello world

speechbrain.utils.distributed.ddp_broadcast(communication_object, src=0)

In DDP mode, this function will broadcast an object to all processes.

Parameters:

• communication_object (Any) – The object to be communicated to all processes. Must be picklable. See docs for torch.distributed.broadcast_object_list()

• src (int) – The rank which holds the object to be communicated.

Return type:

The communication_object passed on rank src.

speechbrain.utils.distributed.ddp_init_group(run_opts)

This function will initialize the ddp group if distributed_launch bool is given in the python command line.

The ddp group will use distributed_backend arg for setting the DDP communication protocol. RANK Unix variable will be used for registering the subprocess to the ddp group.

Parameters:

run_opts (list) – A list of arguments to parse, most often from sys.argv[1:].

Return type:

None