rlightning.utils.ray.remote_class

Ray actor mixin classes for distributed execution.

This module provides mixin classes for Ray actors, including utilities for distributed environment initialization and communication group setup.

class rlightning.utils.ray.remote_class.DistributedMixin

Bases: object

Mixin class for distributed training functionality.

Provides methods for initializing distributed environments and communication groups across Ray actors.

dist_barrier(mode: ParallelMode = None, device_ids: List[int] = None)

Barrier for distributed synchronization.

By default, this barriers on the global process group.

When mode is provided (e.g., ParallelMode.TRAIN_DATA_PARALLEL), it will barrier on that specific communication group. This is important when only a subset of ranks participate in the group; using the global barrier would deadlock.

get_rank() → int

Get the global rank of this process.

Returns:

Global rank.

Raises:

AssertionError – If distributed environment is not initialized.

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

Initialize the distributed environment for this actor.

Sets up environment variables and initializes PyTorch distributed.

Parameters:

• rank – Global rank of this process.

• world_size – Total number of processes.

• local_rank – Local rank within the node.

• local_world_size – Number of processes on this node.

• master_addr – Address of the master node.

• master_port – Port of the master node.

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

• dist_url – URL for distributed initialization.

• timeout – Timeout in seconds for initialization.

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

Initialize a single communication group for the given ranks.

Parameters:

• ranks – List of global ranks to include in the group.

• mode – The parallel mode for this communication group.

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

• use_cpu – If True, use CPU tensors for communication.

Raises:

AssertionError – If distributed environment is not initialized.

class rlightning.utils.ray.remote_class.RayActorMixin

Bases: object

Mixin class providing Ray actor functionality.

This class provides common methods for Ray actors, including GPU management, node identification, and network address retrieval. Should be used as a base class for any class that needs to be a Ray actor.

classmethod as_remote(num_cpus: int | None = None, num_gpus: int | None = None, memory: int | None = None, object_store_memory: int | None = None, resources: Dict[str, float] | None = None) → Type

Create a remote class for Ray Actor initialization.

Parameters:

• num_cpus – Number of CPUs required.

• num_gpus – Number of GPUs required.

• memory – Memory required in bytes.

• object_store_memory – Object store memory required.

• resources – Custom resource requirements.

Returns:

Ray remote class.

init(*args: Any, **kwargs: Any) → None

Initialize the remote class.

Must be implemented in derived classes.

Raises:

NotImplementedError – Always, as this must be overridden.