backends.experimental.utils

Module Contents

Classes

RayStageSpecKeys

String enum of different flags that define keys inside ray_stage_spec.

Functions

execute_setup_on_node	Execute setup on node for a stage.
get_available_cpu_gpu_resources	Get available CPU and GPU resources from Ray.
get_worker_metadata_and_node_id	Get the worker metadata and node id from the runtime context.

API

class backends.experimental.utils.RayStageSpecKeys

Bases: str, enum.Enum

String enum of different flags that define keys inside ray_stage_spec.

Initialization

Initialize self. See help(type(self)) for accurate signature.

IS_ACTOR_STAGE

‘is_actor_stage’

IS_FANOUT_STAGE

‘is_fanout_stage’

IS_LSH_STAGE

‘is_lsh_stage’

IS_RAFT_ACTOR

‘is_raft_actor’

IS_SHUFFLE_STAGE

‘is_shuffle_stage’

backends.experimental.utils.execute_setup_on_node(

stages: list[nemo_curator.stages.base.ProcessingStage],

) → None

Execute setup on node for a stage.

backends.experimental.utils.get_available_cpu_gpu_resources(

init_and_shudown: bool = False,

) → tuple[int, int]

Get available CPU and GPU resources from Ray.

backends.experimental.utils.get_worker_metadata_and_node_id() → tuple[nemo_curator.backends.base.NodeInfo, nemo_curator.backends.base.WorkerMetadata]

Get the worker metadata and node id from the runtime context.