API ReferenceFull Library ReferenceNemo CuratorNemo CuratorBackendsExperimentalRay Actor Pool

nemo_curator.backends.experimental.ray_actor_pool.utils

View as Markdown

Module Contents

Functions

NameDescription
calculate_optimal_actors_for_stageCalculate optimal number of actors for a stage.
create_named_ray_actor_pool_stage_adapterCreate a named RayActorPoolStageAdapter or RayActorPoolRAFTAdapter.

Data

_LARGE_INT

API

nemo_curator.backends.experimental.ray_actor_pool.utils.calculate_optimal_actors_for_stage(
    stage: nemo_curator.stages.base.ProcessingStage,
    num_tasks: int,
    reserved_cpus: float = 0.0,
    reserved_gpus: float = 0.0,
    ignore_head_node: bool = False
) -> int

Calculate optimal number of actors for a stage.

nemo_curator.backends.experimental.ray_actor_pool.utils.create_named_ray_actor_pool_stage_adapter(
    stage: nemo_curator.stages.base.ProcessingStage,
    cls: type[nemo_curator.backends.experimental.ray_actor_pool.adapter.RayActorPoolStageAdapter] | type[nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter]    cls: type[nemo_curator.backends.experimental.ray_actor_pool.adapter.RayActorPoolStageAdapter] | type[nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter]
) -> ray.actor.ActorClass[nemo_curator.backends.experimental.ray_actor_pool.adapter.RayActorPoolStageAdapter | nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter]) -> ray.actor.ActorClass[nemo_curator.backends.experimental.ray_actor_pool.adapter.RayActorPoolStageAdapter | nemo_curator.backends.experimental.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter]

Create a named RayActorPoolStageAdapter or RayActorPoolRAFTAdapter.

This function creates a dynamic subclass of the given adapter class, named after the stage’s class name. This ensures that when Ray calls type(adapter).name, it returns the original stage’s class name rather than ‘RayActorPoolStageAdapter’ or ‘RayActorPoolRAFTAdapter’.

Parameters:

stage
ProcessingStage

ProcessingStage to adapt

cls
type

The adapter class to inherit from

Returns: ActorClass[RayActorPoolStageAdapter | RayActorPoolRAFTAdapter]

A ray.remote decorated class that can be used to create actors

nemo_curator.backends.experimental.ray_actor_pool.utils._LARGE_INT = 2 ** 31 - 1