`nemo_rl.data_plane.mooncake_storage_unit`#

CPU storage-unit placement, and the actors that own Mooncake memory.

Storage units hold all data-plane memory so no compute process does. Both backends place them the same way: storage_unit_placement names the clusters whose nodes host them, and :func:plan_storage_unit_nodes turns that into one Ray node ID per unit, round-robin over those nodes.

mooncake_cpu with storage_unit_segment_size > 0: every other process (trainers, generation, controller) is a client that owns nothing, so a checkpoint save calls only these :class:MooncakeStorageUnit actors.
simple with storage_unit_placement set: TQ’s own SimpleStorageUnits are pinned to the planned nodes instead of TQ’s SPREAD placement group.

Module Contents#

Classes#

MooncakeStorageUnit

Own a Mooncake segment; serve checkpoint and teardown commands.

Functions#

`select_clusters`	The clusters `storage_unit_placement` names (`"all"`: every one).
`storage_node_ids`	Ray node ID per storage unit, round-robin over the clusters’ nodes.
`plan_storage_unit_nodes`	One Ray node ID per storage unit, or None when storage units are off.
`start_storage_units`	Start one MooncakeStorageUnit per entry of `node_ids`.
`stop_storage_units`	Counterpart of :func:`start_storage_units`: run every unit’s teardown.

Data#

`logger`
`_UNIT_STARTUP_TIMEOUT_S`
`_UNIT_CLOSE_TIMEOUT_S`

API#

nemo_rl.data_plane.mooncake_storage_unit.logger#: ‘getLogger(…)’

nemo_rl.data_plane.mooncake_storage_unit._UNIT_STARTUP_TIMEOUT_S#: 600

nemo_rl.data_plane.mooncake_storage_unit._UNIT_CLOSE_TIMEOUT_S#: 60.0

class nemo_rl.data_plane.mooncake_storage_unit.MooncakeStorageUnit(dp_config: nemo_rl.data_plane.DataPlaneConfig)#

Own a Mooncake segment; serve checkpoint and teardown commands.

Initialization

close() → None#: End-of-run teardown; see :func:stop_storage_units.

mooncake_checkpoint( body: dict[str, Any], ) → dict[str, Any] | None#: Run an owner-local checkpoint command; return metadata, never payloads.

nemo_rl.data_plane.mooncake_storage_unit.select_clusters( placement: collections.abc.Sequence[str] | Literal[all], clusters_by_name: collections.abc.Mapping[str, nemo_rl.distributed.virtual_cluster.RayVirtualCluster], ) → list[nemo_rl.distributed.virtual_cluster.RayVirtualCluster]#

The clusters storage_unit_placement names ("all": every one).

Names are the keys of clusters_by_name (train, inference, teacher:<name>). A colocated run registers one cluster under two names, so excluding one of them excludes nothing; that is logged.

nemo_rl.data_plane.mooncake_storage_unit.storage_node_ids( clusters: collections.abc.Sequence[nemo_rl.distributed.virtual_cluster.RayVirtualCluster], count: int | None, ) → list[str]#

Ray node ID per storage unit, round-robin over the clusters’ nodes.

count None means 2 per selected node.

nemo_rl.data_plane.mooncake_storage_unit.plan_storage_unit_nodes( dp_config: nemo_rl.data_plane.DataPlaneConfig, clusters_by_name: collections.abc.Mapping[str, nemo_rl.distributed.virtual_cluster.RayVirtualCluster], ) → list[str] | None#: One Ray node ID per storage unit, or None when storage units are off.

nemo_rl.data_plane.mooncake_storage_unit.start_storage_units( dp_config: nemo_rl.data_plane.DataPlaneConfig, node_ids: list[str] | None, ) → tuple[Any, ...]#

Start one MooncakeStorageUnit per entry of node_ids.

Returns () for any other backend or when node_ids is None.

async nemo_rl.data_plane.mooncake_storage_unit.stop_storage_units( units: collections.abc.Sequence[Any], ) → None#

Counterpart of :func:start_storage_units: run every unit’s teardown.

Ray stops the units with SIGTERM, and Mooncake’s SIGTERM handler exits without running Python exit hooks, so a unit cleans up only when asked. Bounded and never raises: teardown must not hang or mask the run’s own outcome (a unit still busy with a checkpoint command is skipped).

nemo_rl.data_plane.mooncake_storage_unit#