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.
`OncePerNodeContext`	Context manager to ensure code runs only once per node.

Functions:

`ddp_all_reduce`	In DDP mode, this function will perform an all_reduce operation with the specified torch operator.
`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.
`ddp_prevent_block`	Prevent blocking because only one or partial threads running.
`get_local_rank`	Get the local rank of the current process on the current node.
`get_rank`	Get the rank of the current process.
`if_main_process`	Returns whether the current process is the main process.
`infer_device`	Make a basic guess about intended running device based on availability and distributed environment variable 'LOCAL_RANK'
`is_distributed_initialized`	Returns whether the current system is distributed.
`is_local_rank_zero`	Returns whether the current process has local rank of 0.
`main_process_only`	Function decorator to ensure the function runs only on the main process.
`once_per_node`	Function decorator to ensure the function runs only once per node.
`rank_prefixed_message`	Prefix a message with the rank of the process.
`run_on_main`	Runs a function with DPP (multi-gpu) support.
`run_once_per_node`	Runs a function with DPP (multi-gpu) support.

Reference

speechbrain.utils.distributed.rank_prefixed_message(message: str) → str[source]

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[source]

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.get_local_rank() → int | None[source]

Get the local rank of the current process on the current node.

Returns:: The local rank of the current process, or None if the local rank could not be determined.
Return type:: int or None

speechbrain.utils.distributed.infer_device() → str[source]: Make a basic guess about intended running device based on availability and distributed environment variable ‘LOCAL_RANK’

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)[source]

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)

Returns:

On all processes (the value that func returned, when it ran on the main)
process.

speechbrain.utils.distributed.run_once_per_node(func, args=None, kwargs=None, post_func=None, post_args=None, post_kwargs=None, run_post_on_all=False)[source]

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

The provided function func is only run once on each node, while other processes block to wait for the function execution to finish. This is useful for things such as saving a file to the disk on each separate node (i.e. the filesystems are separate). In addition, a second function can be specified to be run on other processes after the first function completes, for example, loading a file that was created on each node.

Parameters:

func (callable) – Function to be run once on each node.
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. By default, post_func is not run on the process that ran func.
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_all (bool) – Whether to run post_func on all processes, including the process that ran func.

Returns:

If post_func is provided, returns the result on all processes where post_func is run.
If run_post_on_all is False or post_func is not provided, returns the result of func on the processes where it is run.
If post_func is not provided, returns None on processes where func was not called.

Example

>>> tmpfile = getfixture("tmpdir") / "example.pt"
>>> # Return tensor so we don't have to load it on the saving process
>>> def save_and_return(file, tensor):
...     torch.save(tensor, file)
...     return tensor
>>> # Load tensor on non-saving processes
>>> def load_tensor(file):
...     return torch.load(file)
>>> # Save on node-primary processes, load on others
>>> example_tensor = torch.ones(5)
>>> loaded_tensor = run_once_per_node(
...     func=save_and_return,
...     args=[tmpfile, example_tensor],
...     post_func=load_tensor,
...     post_args=[tmpfile],
...     run_post_on_all=False,
... )
>>> # We should get the same result on all processes
>>> loaded_tensor
tensor([1., 1., 1., 1., 1.])

speechbrain.utils.distributed.is_distributed_initialized() → bool[source]: Returns whether the current system is distributed.

speechbrain.utils.distributed.if_main_process() → bool[source]: Returns whether the current process is the main process.

speechbrain.utils.distributed.is_local_rank_zero() → bool[source]: Returns whether the current process has local rank of 0.

class speechbrain.utils.distributed.MainProcessContext[source]

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__()[source]: Enter the context. Increase the counter.

__exit__(exc_type, exc_value, traceback)[source]: Exit the context. Decrease the counter.

class speechbrain.utils.distributed.OncePerNodeContext[source]

Bases: object

Context manager to ensure code runs only once per node. This is useful to make sure that NODE_ONCE_ONLY global variable is decreased even if there’s an exception raised inside of the once_per_node_wrapped_fn function.

__enter__()[source]: Enter the context. Increase the counter.

__exit__(exc_type, exc_value, traceback)[source]: Exit the context. Decrease the counter.

speechbrain.utils.distributed.main_process_only(function)[source]: 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. The function will return the result computed on the main process to all processes.

speechbrain.utils.distributed.once_per_node(function)[source]

Function decorator to ensure the function runs only once per node. This is useful for things like saving to the filesystem where you only want it to happen on a single process on each node.

Unlike main_process_only, no broadcasting is done. Instead, processes with local_rank == 0 keep their own result, all other processes return None.

speechbrain.utils.distributed.ddp_prevent_block()[source]: Prevent blocking because only one or partial threads running.

speechbrain.utils.distributed.ddp_barrier()[source]

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)[source]

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_all_reduce(communication_object, reduce_op)[source]

In DDP mode, this function will perform an all_reduce operation with the specified torch operator.

See: https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce

Parameters:

communication_object (Any) – The object to be reduced across processes.
reduce_op (torch.distributed.ReduceOp) – The operation to perform. E.g. include torch.distributed.ReduceOp.AVG or torch.distributed.ReduceOp.SUM. See the Torch documentation for more.

Return type:

The communication_object once reduced (or itself if DDP not initialised)

speechbrain.utils.distributed.ddp_init_group(run_opts)[source]

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