nemo_curator.backends.ray_actor_pool.utils

Module Contents

Classes

Name	Description
`_NoResourcesAvailableError`	-

Functions

Name	Description
`calculate_optimal_actors_for_resources`	Calculate an actor count from a supplied CPU/GPU availability snapshot.
`calculate_optimal_actors_for_stage`	Calculate optimal number of actors for a stage.
`calculate_optimal_actors_for_stage_with_wait`	Wait for enough resources to create the baseline-sized actor pool.
`create_named_ray_actor_pool_stage_adapter`	Create a named RayActorPoolStageAdapter or RayActorPoolRAFTAdapter.
`get_available_actor_pool_resources`	Return currently available CPU/GPU resources after executor reservations.
`update_resource_baseline`	Update a CPU/GPU high-water mark from current availability.

Data

_LARGE_INT

API

class nemo_curator.backends.ray_actor_pool.utils._NoResourcesAvailableError()

Bases: ValueError

nemo_curator.backends.ray_actor_pool.utils.calculate_optimal_actors_for_resources(
    stage: nemo_curator.stages.base.ProcessingStage,
    num_tasks: int,
    available_resources: tuple[float, float],
    ignore_head_node: bool = False
) -> int

Calculate an actor count from a supplied CPU/GPU availability snapshot.

nemo_curator.backends.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.ray_actor_pool.utils.calculate_optimal_actors_for_stage_with_wait(
    stage: nemo_curator.stages.base.ProcessingStage,
    num_tasks: int,
    resource_baseline: tuple[float, float],
    reserved_cpus: float = 0.0,
    reserved_gpus: float = 0.0,
    ignore_head_node: bool = False,
    timeout: float = 5.0,
    interval: float = 0.2
) -> int

Wait for enough resources to create the baseline-sized actor pool.

nemo_curator.backends.ray_actor_pool.utils.create_named_ray_actor_pool_stage_adapter(
    stage: nemo_curator.stages.base.ProcessingStage,
    cls: type[nemo_curator.backends.ray_actor_pool.adapter.RayActorPoolStageAdapter] | type[nemo_curator.backends.ray_actor_pool.raft_adapter.RayActorPoolRAFTAdapter]
) -> ray.actor.ActorClass[nemo_curator.backends.ray_actor_pool.adapter.RayActorPoolStageAdapter | nemo_curator.backends.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 — 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.ray_actor_pool.utils.get_available_actor_pool_resources(
    reserved_cpus: float = 0.0,
    reserved_gpus: float = 0.0,
    ignore_head_node: bool = False
) -> tuple[float, float]

Return currently available CPU/GPU resources after executor reservations.

nemo_curator.backends.ray_actor_pool.utils.update_resource_baseline(
    resource_baseline: tuple[float, float],
    reserved_cpus: float = 0.0,
    reserved_gpus: float = 0.0,
    ignore_head_node: bool = False
) -> tuple[float, float]

Update a CPU/GPU high-water mark from current availability.

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