rlightning.utils.placement.resource_pool

Resource Pool Planner for global resource management.

This module provides functionality for: 1. Discovering and managing cluster node resources 2. Planning resource pools based on component scheduling requirements 3. Supporting different allocation strategies (disaggregate, colocate) 4. Tracking fine-grained component-to-resource mappings

class rlightning.utils.placement.resource_pool.ComponentAllocation(component_type: str, node_to_gpu_indices: Dict[str, List[str]], total_gpus: int, worker_count: int)

Bases: object

Represents resource allocation for a specific component type.

component_type

Type of component (train, eval, env, buffer)

Type:

str

gpu_indices

List of GPU index ranges, e.g., [“0-7”, “8-15”]

total_gpus

Total number of GPUs allocated

Type:

int

worker_count

Number of workers for this component

Type:

int

component_type: str

get_node_index_string(node_id: str) → str

Get index string for a specific node_id.

node_to_gpu_indices: Dict[str, List[str]]

to_dict() → Dict[str, Any]

Convert to dictionary representation.

to_index_string() → str

Convert GPU indices to a compact string format like ‘0-7, 8-15’.

total_gpus: int

worker_count: int

class rlightning.utils.placement.resource_pool.NodeResource(node_id: str, ip: str, total_cpus: int, total_gpus: int, allocations: Dict[str, ~typing.List[~typing.Tuple[int, int]]]=<factory>, gpu_cursor: int = 0, max_allocated_gpus: int = 0)

Bases: object

Represents a node’s resources.

This class is used in two contexts: - Cluster discovery: total_* reflects node total capacity; available_* is computed. - ResourcePool membership: total_* represents the node’s capacity,

and allocations records per-component GPU index ranges within the node’s index space.

Resource tracking uses gpu_cursor to track the next available GPU index. available_gpus is computed as total_gpus - gpu_cursor.

allocate(gpus: int | List[int], component_types: List[str] | None = None, consume: bool = True) → None

Allocate resources directly on this node.

Modifies self.allocations and advances gpu_cursor.

Parameters:

• gpus – GPU units to allocate. Can be int (same for all) or list (per-component).

• component_types – Components to record allocations for. If None, just advances cursor.

• consume – If True, consume the GPUs from the node.

allocations: Dict[str, List[Tuple[int, int]]]

property available_gpus: int

Remaining GPUs available for allocation.

property component_types: List[str]

Get list of component types that have allocations on this node.

copy() → NodeResource

Create a deep copy of this node resource.

gpu_cursor: int = 0

has_resources(cpus: int = 0, gpus: int = 0) → bool

Check if node has sufficient available resources.

ip: str

property is_empty: bool

Check if node has no more allocatable resources.

max_allocated_gpus: int = 0

node_id: str

total_cpus: int

total_gpus: int

class rlightning.utils.placement.resource_pool.ResourcePool(name: str, nodes: ~typing.List[~rlightning.utils.placement.resource_pool.NodeResource], _component_types: ~typing.List[str] | None = None, component_allocations: ~typing.Dict[str, ~rlightning.utils.placement.resource_pool.ComponentAllocation] = <factory>)

Bases: object

A resource pool containing allocated nodes for specific components.

This class tracks: - Which nodes belong to this pool - Which component types use this pool - Fine-grained GPU index allocations per component per node

component_types is auto-inferred from node allocations if not provided. If ‘train’ exists in allocations, ‘buffer’ is automatically added.

component_allocations: Dict[str, ComponentAllocation]

property component_types: List[str]

Get component types for this pool.

Auto-infers from node allocations if not explicitly set. Auto-adds ‘buffer’ if ‘train’ exists.

classmethod from_yaml_dict(pool_cfg: Dict[str, Any], cluster_nodes: Dict[str, NodeResource], *, used_node_ids: set[str] | None = None) → ResourcePool

Build a ResourcePool from a YAML pool dict.

Supported input fields:

• name: str

• num_node: int

• num_gpus: int (per-node) OR List[int] (per-node gpus list)

• optional node_ids: List[str] to pin pools to specific nodes

• component keys: train/eval/env/buffer/… with values like “0-7, 8-15”

component_types is auto-inferred from allocations if not explicitly listed. ‘buffer’ is auto-added if ‘train’ exists.

get_component_indices(component_type: str) → str

Get the GPU index string for a component type across all nodes.

Returns:

Index string like “0-7” or “0-7, 8-15” for multi-node.

get_component_node_count(component_type: str) → int

Get the number of nodes that have the component type.

name: str

property node_ids: List[str]

List of node IDs in this pool.

nodes: List[NodeResource]

property num_nodes: int

Number of nodes in this pool.

to_dict() → Dict[str, Any]

Convert to dictionary for serialization.

to_yaml_dict() → Dict[str, Any]

Convert to YAML-friendly dictionary format.

Output format:

name: “train_pool” num_node: 1 num_gpus: 8 # per-node allocated gpus, NOT total train: “0-7”

property total_cpus: int

Total available CPUs in this pool.

property total_gpus: int

Total available GPUs in this pool.

class rlightning.utils.placement.resource_pool.ResourcePoolPlanner(scheduling: ComponentScheduling, cluster_info: Dict[str, Any] | None = None)

Bases: object

Plans and manages resource pools for component placement.

This planner: 1. Discovers cluster resources 2. Validates scheduling requirements against available resources 3. Allocates resources to pools based on the placement strategy 4. Tracks fine-grained component-to-GPU mappings

discover_cluster_resources() → Dict[str, NodeResource]

Discover and cache cluster node resources.

Returns:

Dictionary mapping node_id to NodeResource.

get_cluster_info() → Dict[str, Any]

Get raw cluster information.

get_component_node_count(component_type: str) → int

Get the number of nodes that have the component type.

get_node_resources() → Dict[str, NodeResource]

Get node resources dictionary.

get_pool_for_component(component_type: str) → ResourcePool | None

Get the resource pool that contains a specific component type.

get_resource_pools(pool_name: str = None) → Dict[str, ResourcePool]

Get the planned resource pools.

load_manual_resource_pools(resource_pool_cfg: List[Dict[str, Any]]) → Dict[str, ResourcePool]

Load resource pools from a manual YAML config list.

plan_resource_pools(strategy: str = 'disaggregate', env_strategy: str = 'default') → Dict[str, ResourcePool]

Plan resource pools based on placement strategy.

Parameters:

• strategy – The placement strategy to use.

• env_strategy – The environment placement strategy.

Returns:

Dictionary mapping pool name to ResourcePool.

resource_summary: Dict[str, Any]

summary() → Dict[str, Any]

Get a summary of the resource planning.

to_yaml_config() → Dict[str, Any]

Generate a YAML-compatible configuration representing the resource pools.

Output format:

resource_pool:

• name: “train_pool” num_node: 1 num_gpus: 8 # per-node GPUs train: “0-7”

• name: “rollout_pool” num_node: 2 num_gpus: 8 # per-node GPUs eval: “0-7, 8-15” env: “0-7, 8-15”

validate_scheduling(strategy: str = 'disaggregate', env_strategy: str = 'default') → Tuple[bool, str]

Validate that cluster resources can satisfy scheduling requirements.

Parameters:

• strategy – The placement strategy to use.

• env_strategy – The environment strategy to use.

Returns:

Tuple of (is_valid, error_message).