`nemo_rl.distributed.virtual_cluster`#

Module Contents#

Classes#

`ClusterConfig`
`PY_EXECUTABLES`
`RayVirtualCluster`	Creates a virtual distributed cluster using Ray placement groups.

Functions#

`_get_node_ip_and_free_port`
`_get_node_ip_local`
`_bind_socket_in_range`	Try to bind sock to a random port in [port_range_low, port_range_high).
`_get_free_port_local`
`init_ray`	Initialise Ray.
`_get_gpu_id_info`	Return (gpu_id, nvlink_domain, topo_rank) for the current worker’s bundle.
`get_reordered_bundle`	Return bundle indices and GPU IDs ordered by physical topology.
`get_ray_cluster_topology`	Query all alive Ray nodes for their NVLink domain and topo_rank.
`select_segment_nodes`	Partition Ray node IDs into segment-aligned selected nodes and remainder.
`prepare_segment_topology`	Compute node resource constraints for topology-aware cluster placement.
`_sort_bundle_indices_by_topology`	Compute topology-aware sort order for bundle indices.

Data#

`logger`
`dir_path`
`git_root`
`DEFAULT_GENERATION_PORT_RANGE_LOW`
`DEFAULT_GENERATION_PORT_RANGE_HIGH`
`DEFAULT_GYM_PORT_RANGE_LOW`
`DEFAULT_GYM_PORT_RANGE_HIGH`
`DEFAULT_VLLM_PORT_RANGE_LOW`
`DEFAULT_VLLM_PORTS_PER_ENGINE`
`DEFAULT_MASTER_PORT_RANGE_LOW`
`DEFAULT_MASTER_PORT_RANGE_HIGH`
`NVLINK_DOMAIN_PREFIX`	Ray resource key prefix for the NVLink domain. Each node registers one resource `nvlink_domain_<ClusterUUID>: 1` where ClusterUUID is parsed directly from `nvidia-smi -q` output by ray.sub. Nodes sharing the same key belong to the same NVLink switch fabric (e.g. one GB200 NVL72 rack).
`TOPO_RANK_KEY`	Ray resource key for the SLURM topological rank. Derived from `SLURM_TOPOLOGY_ADDR` (when `SLURM_TOPOLOGY_ADDR_PATTERN=block.node`), falling back to `SLURM_PROCID + 2` on worker nodes (head node is pinned to `1`), then to hostname digits when SLURM is unavailable. Values are always `>= 1` so that Ray does not drop the custom resource (Ray drops value-0 custom resources). Used to sort nodes within and across NVLink domains so rank assignment follows physical topology.
`NVLINK_DOMAIN_UNKNOWN`	Sentinel returned when no NVLink domain info is available for a node.
`TOPO_RANK_UNKNOWN`	Sentinel returned when no topological rank is available for a node.

API#

nemo_rl.distributed.virtual_cluster.logger#: ‘getLogger(…)’

class nemo_rl.distributed.virtual_cluster.ClusterConfig#

Bases: typing.TypedDict

gpus_per_node: int#: None

num_nodes: int#: None

master_port_range_low: NotRequired[int]#: None

master_port_range_high: NotRequired[int]#: None

segment_size: NotRequired[int | None]#: None

nemo_rl.distributed.virtual_cluster.dir_path#: ‘dirname(…)’

nemo_rl.distributed.virtual_cluster.git_root#: ‘abspath(…)’

class nemo_rl.distributed.virtual_cluster.PY_EXECUTABLES#

SYSTEM#: None

BASE#: None

VLLM#: None

FSDP#: None

AUTOMODEL#: None

MCORE#: None

NEMO_GYM#: None

SGLANG#: None

nemo_rl.distributed.virtual_cluster.DEFAULT_GENERATION_PORT_RANGE_LOW#: 3000

nemo_rl.distributed.virtual_cluster.DEFAULT_GENERATION_PORT_RANGE_HIGH#: 4999

nemo_rl.distributed.virtual_cluster.DEFAULT_GYM_PORT_RANGE_LOW#: 5000

nemo_rl.distributed.virtual_cluster.DEFAULT_GYM_PORT_RANGE_HIGH#: 5999

nemo_rl.distributed.virtual_cluster.DEFAULT_VLLM_PORT_RANGE_LOW#: 7000

nemo_rl.distributed.virtual_cluster.DEFAULT_VLLM_PORTS_PER_ENGINE#: 100

nemo_rl.distributed.virtual_cluster.DEFAULT_MASTER_PORT_RANGE_LOW#: 1400

nemo_rl.distributed.virtual_cluster.DEFAULT_MASTER_PORT_RANGE_HIGH#: 1999

nemo_rl.distributed.virtual_cluster.NVLINK_DOMAIN_PREFIX#

‘nvlink_domain_’

Ray resource key prefix for the NVLink domain. Each node registers one resource nvlink_domain_<ClusterUUID>: 1 where ClusterUUID is parsed directly from nvidia-smi -q output by ray.sub. Nodes sharing the same key belong to the same NVLink switch fabric (e.g. one GB200 NVL72 rack).

nemo_rl.distributed.virtual_cluster.TOPO_RANK_KEY#

‘topo_rank’

Ray resource key for the SLURM topological rank. Derived from SLURM_TOPOLOGY_ADDR (when SLURM_TOPOLOGY_ADDR_PATTERN=block.node), falling back to SLURM_PROCID + 2 on worker nodes (head node is pinned to 1), then to hostname digits when SLURM is unavailable. Values are always >= 1 so that Ray does not drop the custom resource (Ray drops value-0 custom resources). Used to sort nodes within and across NVLink domains so rank assignment follows physical topology.

nemo_rl.distributed.virtual_cluster.NVLINK_DOMAIN_UNKNOWN#

‘unknown’

Sentinel returned when no NVLink domain info is available for a node.

nemo_rl.distributed.virtual_cluster.TOPO_RANK_UNKNOWN: int#

None

Sentinel returned when no topological rank is available for a node.

nemo_rl.distributed.virtual_cluster._get_node_ip_and_free_port( port_range_low: int = DEFAULT_MASTER_PORT_RANGE_LOW, port_range_high: int = DEFAULT_MASTER_PORT_RANGE_HIGH, ) → tuple[str, int]#

nemo_rl.distributed.virtual_cluster._get_node_ip_local() → str#

nemo_rl.distributed.virtual_cluster._bind_socket_in_range( sock: socket.socket, port_range_low: int, port_range_high: int, max_retries: int = 50, ) → int#

Try to bind sock to a random port in [port_range_low, port_range_high).

Raises RuntimeError after max_retries failed attempts.

nemo_rl.distributed.virtual_cluster._get_free_port_local( port_range_low: int = DEFAULT_MASTER_PORT_RANGE_LOW, port_range_high: int = DEFAULT_MASTER_PORT_RANGE_HIGH, ) → int#

nemo_rl.distributed.virtual_cluster.init_ray(log_dir: Optional[str] = None) → None#

Initialise Ray.

Try to attach to an existing local cluster. If that cluster uses the same CUDA_VISIBLE_DEVICES or Slurm managed tag we will reuse it. Otherwise, we will detach and start a fresh local cluster.

Parameters:: log_dir – Optional directory to store Ray logs and temp files.

nemo_rl.distributed.virtual_cluster._get_gpu_id_info() → tuple[int, str, int]#

Return (gpu_id, nvlink_domain, topo_rank) for the current worker’s bundle.

Reads custom resources set by ray.sub (see NVLINK_DOMAIN_PREFIX / TOPO_RANK_KEY).

nemo_rl.distributed.virtual_cluster.get_reordered_bundle( pg: ray.util.placement_group.PlacementGroup, segment_size: int | None = None, gpus_per_node: int | None = None, ) → tuple[list[int], list[int], tuple[str, ...]]#

Return bundle indices and GPU IDs ordered by physical topology.

Spins up one short-lived task per bundle (_get_gpu_id_info) to discover each bundle’s physical GPU ID, NVLink domain, and topo_rank, then orders bundles via _sort_bundle_indices_by_topology:

No topology info available -> sort by (node_id, gpu_id).
Topology info available -> sort by (domain_min_topo_rank, topo_rank, gpu_id).
segment_size set -> additionally drop bundles from NVLink domains that cannot form a complete segment_size-node segment.

Parameters:

pg – Placement group whose bundles to reorder.
segment_size – Nodes per NVLink domain segment for topology-aware alignment; None disables segment filtering.
gpus_per_node – Required when segment_size is set.

Returns:

(reordered_bundle_indices, reordered_gpu_ids, nvlink_domain_per_bundle_index) where nvlink_domain_per_bundle_index is indexed by original bundle index.

exception nemo_rl.distributed.virtual_cluster.ResourceInsufficientError#

Bases: Exception

Exception raised when the cluster does not have enough resources to satisfy the requested configuration.

Initialization

Initialize self. See help(type(self)) for accurate signature.

nemo_rl.distributed.virtual_cluster.get_ray_cluster_topology() → dict[str, tuple[str, int]]#

Query all alive Ray nodes for their NVLink domain and topo_rank.

Returns:: Dict mapping node_id -> (nvlink_domain, topo_rank). nvlink_domain is NVLINK_DOMAIN_UNKNOWN and topo_rank is TOPO_RANK_UNKNOWN if topology info is unavailable.

nemo_rl.distributed.virtual_cluster.select_segment_nodes( topology: dict[str, tuple[str, int]], segment_size: int, num_nodes: int, ) → tuple[list[str], list[str]]#

Partition Ray node IDs into segment-aligned selected nodes and remainder.

Greedily selects complete segments (segment_size nodes) from each NVLink domain, sorted by topological order, until num_nodes is reached.

Parameters:

topology – Dict mapping node_id -> (nvlink_domain, topo_rank) from get_ray_cluster_topology().
segment_size – Number of nodes per NVLink domain segment.
num_nodes – Total number of nodes to select.

Returns:

Selected nodes are in topological order.

Return type:

(selected_node_ids, remaining_node_ids)

Raises:

ValueError – If segment_size does not evenly divide num_nodes.
ResourceInsufficientError – If not enough complete segments can be formed.

nemo_rl.distributed.virtual_cluster.prepare_segment_topology( segment_size: int | None, num_nodes: int, *, topology: dict[str, tuple[str, int]] | None = None, role: str = 'training', ) → tuple[list[dict[str, float]] | None, list[str], dict[str, tuple[str, int]]]#

Compute node resource constraints for topology-aware cluster placement.

Fetches cluster topology if not provided, selects segment-aligned nodes, and returns per-node domain constraints ready for RayVirtualCluster.

Parameters:

segment_size – Nodes per NVLink domain segment. None disables topology logic.
num_nodes – Number of nodes to select.
topology – Pre-fetched topology dict; fetched automatically when None.
role – Label used in progress messages (e.g. "training", "inference").

Returns:

(node_resource_constraints, remaining_node_ids, topology)

node_resource_constraints: per-node domain-pinning dicts for RayVirtualCluster, or None when segment_size is None or no NVLink domain info is available.
remaining_node_ids: node IDs not selected; pass the corresponding sub-topology to a follow-up call to allocate an inference cluster.
topology: the topology dict used (empty dict when segment_size is None, for safe sub-topology slicing by callers).

nemo_rl.distributed.virtual_cluster._sort_bundle_indices_by_topology( bundle_data: list[tuple[int, str, int, str]], segment_size: int | None = None, gpus_per_node: int | None = None, ) → list[int]#

Compute topology-aware sort order for bundle indices.

When topology information is available: sort by (domain_min_topo_rank, topo_rank, gpu_id). When segment_size is set: additionally validate that each NVLink domain contributes complete segments (segment_size nodes), discarding bundles from incomplete domains. Else: sort by (node_id, gpu_id).

Parameters:

bundle_data – For each bundle i, (gpu_id, nvlink_domain, topo_rank, node_id).
segment_size – If set, number of nodes per NVLink domain segment. Bundles from domains with fewer than segment_size nodes are excluded.
gpus_per_node – Required when segment_size is set. Number of GPUs per node.

Returns:

List of bundle indices in sorted order.

Raises:

ValueError – If segment_size is set but gpus_per_node is not.

class nemo_rl.distributed.virtual_cluster.RayVirtualCluster( bundle_ct_per_node_list: list[int], use_gpus: bool = True, max_colocated_worker_groups: int = 1, num_gpus_per_node: int = 8, name: str = '', placement_group_strategy: str = 'SPREAD', port_range_low: Optional[int] = None, port_range_high: Optional[int] = None, segment_size: int | None = None, node_resource_constraints: list[dict[str, float]] | None = None, )#

Creates a virtual distributed cluster using Ray placement groups.

This class simplifies distributed training setup by:

Creating placement groups that represent logical compute nodes
Allocating GPU and CPU resources for distributed workers
Managing communication between distributed processes
Bundle: A resource allocation unit (ex: 4 GPUs on a single node)
Worker: A process that performs computation (model training/inference)
Node: A physical or virtual machine containing multiple bundles

Initialization

Initialize a virtual cluster using Ray placement groups.

Parameters:

bundle_ct_per_node_list – List specifying GPU bundles per node (e.g., [2,2] creates 2 nodes with 2 GPU bundles each)
use_gpus – Whether to allocate GPU resources
max_colocated_worker_groups – Maximum number of worker groups that can be colocated
num_gpus_per_node – Number of GPUs per node
name – Name prefix for placement groups
placement_group_strategy – Ray placement group strategy (“STRICT_PACK”, “PACK”, or “SPREAD”)
port_range_low – Lower bound (inclusive) of the port range for master address allocation. Falls back to DEFAULT_MASTER_PORT_RANGE_LOW if None.
port_range_high – Upper bound (exclusive) of the port range for master address allocation. Falls back to DEFAULT_MASTER_PORT_RANGE_HIGH if None.
segment_size – Nodes per NVLink domain segment for topology-aware alignment. When set, _sort_bundle_indices_by_topology trims incomplete domain segments.
node_resource_constraints – Per-logical-node extra Ray resource requirements. Length must match bundle_ct_per_node_list. Each dict is merged into every bundle spec for that node, pinning it to a physical domain. Built from NVLink domain resources injected by ray.sub. Example: [{“nvlink_domain_”: 0.001}] * 16 pins 16 nodes to a single NVLink domain.

_init_placement_groups( strategy: str | None = None, use_unified_pg: bool = False, ) → list[ray.util.placement_group.PlacementGroup]#

Creates placement groups based on whether cross-node model parallelism is needed.

Parameters:

strategy – Ray placement group strategy (defaults to self.placement_group_strategy)
use_unified_pg – If True, create a single unified placement group. If False, create per-node placement groups.

Returns:

List of placement groups

_create_placement_groups_internal( strategy: str, use_unified_pg: bool = False, ) → list[ray.util.placement_group.PlacementGroup]#: Internal method to create placement groups without retry logic.

get_placement_groups() → list[ray.util.placement_group.PlacementGroup]#

world_size() → int#

node_count() → int#

get_available_address_and_port( pg_idx: int, bundle_idx: int, ) → tuple[str, int]#

Gets an available address and port for the given placement group index and bundle index.

Returns:: Tuple of (address, port)

get_available_addresses_and_ports_batch( pg_bundle_pairs: list[tuple[int, int]], batch_size: int = 256, ) → list[tuple[str, int]]#

Discovers available addresses and ports for multiple bundles in parallel.

Fires all remote tasks up front, then collects results in batches via ray.wait() to avoid putting too many objects into the Ray object store at once. See https://docs.ray.io/en/latest/ray-core/patterns/ray-get-too-many-objects.html

Parameters:

pg_bundle_pairs – List of (pg_idx, bundle_idx) pairs.
batch_size – Maximum number of ready futures to fetch at once.

Returns:

List of (address, port) tuples in the same order as pg_bundle_pairs.

get_master_address_and_port() → tuple[str, int]#

Gets the master address and port for the distributed training setup.

Each call returns a unique port that has not been returned by previous calls on this cluster instance. This prevents NCCL process-group collisions when multiple worker groups (e.g. policy and value) share the same cluster and node.

Returns:: Tuple of (address, port)

_get_sorted_bundle_indices() → Optional[list[int]]#

Gets the sorted bundle indices for the placement groups.

Returns:: List of bundle indices in sorted order.

shutdown() → bool#

Cleans up and releases all resources associated with this virtual cluster.

This includes removing all placement groups and resetting the internal state.

This method is idempotent and can be safely called multiple times.

__del__() → None#

Shutsdown the virtual cluster when the object is deleted or is garbage collected.

This is an extra safety net in case the user forgets to call shutdown and the pointer to the cluster is lost due to leaving a function scope. It’s always recommended that the user calls shutdown().

nemo_rl.distributed.virtual_cluster#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.distributed.virtual_cluster`#