nemo_curator.backends.experimental.ray_actor_pool.raft_adapter

Module Contents

Classes

Name	Description
`RayActorPoolRAFTAdapter`	RAFT Actor adapter for Ray Actor Pool backend.

API

class nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter(
    stage: nemo_curator.stages.base.ProcessingStage,
    index: int,
    pool_size: int,
    session_id: bytes,
    actor_name_prefix: str = 'RAFT'
)

Bases: BaseStageAdapter

RAFT Actor adapter for Ray Actor Pool backend.

This adapter extends RayActorPoolStageAdapter and adds RAFT capabilities to enable distributed processing with RAFT communication.

_batch_size

= self.stage.batch_size

_is_root

= not index

_name

= f'{self._actor_name_prefix}Actor-{self._index}'

root_unique_id

= self.unique_id if self._is_root else None

unique_id

= self.cb.uniqueId

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter._setup_nccl() -> None

Setup NCCL communicator.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter._setup_raft() -> None

Setup RAFT.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter.broadcast_root_unique_id() -> None

Broadcast the root unique ID to all actors.

This method should only be called by the root actor.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter.get_batch_size() -> int

Get the batch size for this stage.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter.set_root_unique_id(
    root_unique_id: int
) -> None

Set the root unique ID.

Parameters

root_unique_id : int The root unique ID.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter.setup(
    worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

Setup the RAFT actor.

This method should be called after the root unique ID has been broadcast.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter.setup_on_node() -> None

Setup method for Ray actors.

Note: This method is not used in the current implementation since we use the Ray Data pattern of calling setup_on_node before actor creation.

nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter.teardown() -> None