rlightning.utils.distributed.comm_context

Communication context for distributed training.

This module provides the CommContext singleton class for managing PyTorch distributed communication groups and process ranks.

class rlightning.utils.distributed.comm_context.CommContext(*args: Any, **kwargs: Any)

Bases: object

Singleton communication context for PyTorch distributed.

Manages process groups, ranks, and world sizes for various communication modes in distributed training.

_local_ranks

Mapping from CommMode to local rank.

_global_ranks

Mapping from CommMode to global rank.

_world_sizes

Mapping from CommMode to world size.

_ranks_in_group

Mapping from CommMode to list of ranks.

_groups

Mapping from CommMode to ProcessGroup.

destroy() → None

Destroy the distributed environment and clear all groups.

get_global_rank() → int

Get the global rank of this process.

Returns:

Global rank.

get_group(comm_mode: CommMode)

Get the process group for a communication mode.

Parameters:

comm_mode – The communication mode.

Returns:

The corresponding ProcessGroup.

get_in_node_rank() → int

Get the local rank within the current node.

Returns:

Local rank from LOCAL_RANK environment variable.

get_inter_node_process_group()

Get the inter-node process group.

Returns:

Inter-node ProcessGroup.

get_intra_node_process_group()

Get the intra-node process group.

Returns:

Intra-node ProcessGroup.

get_local_rank(comm_mode: CommMode) → int

Get the local rank for a communication mode.

Parameters:

comm_mode – The communication mode.

Returns:

Local rank within the group.

get_local_world_size() → int

Get the world size within the current node.

Returns:

Local world size from LOCAL_WORLD_SIZE environment variable.

get_ranks_in_group(comm_mode: CommMode) → List[int]

Get all ranks in a communication group.

Parameters:

comm_mode – The communication mode.

Returns:

List of global ranks in the group.

get_world_size(comm_mode: CommMode) → int

Get the world size for a communication mode.

Parameters:

comm_mode – The communication mode.

Returns:

World size for the mode.

init_distributed_env(world_size: int | None = None, rank: int | None = None, backend: str = 'nccl', dist_url: str = 'env://', timeout: int = 1800) → None

Initialize the PyTorch distributed process group.

Parameters:

• world_size – Global world size.

• rank – Global rank of this process.

• backend – Communication backend (‘nccl’ or ‘gloo’).

• dist_url – Initialization method URL.

• timeout – Timeout in seconds for initialization.

init_group(ranks: List[int], mode: ParallelMode | None = None, backend: str = 'nccl', use_cpu: bool = False) → None

Initialize a communication group for the given mode.

Parameters:

• ranks – List of global ranks to include.

• mode – The parallel mode for this group.

• backend – Communication backend.

• use_cpu – If True, use CPU tensors.

Raises:

NotImplementedError – If the parallel mode is not supported.

is_initialized() → bool

Check if the global communication group is initialized.

Returns:

True if initialized, False otherwise.

is_main_rank() → bool

Check if this is the main (rank 0) process.

Returns:

True if global rank is 0.