`nemo_rl.algorithms.metric_utils`#

Module Contents#

Classes#

SetupTimingMetrics

Driver-side per-phase timings collected during setup.

Functions#

print_setup_timing_summary

Print the setup-phase summary block.

Data#

`SETUP_TIMING_PREFIX`
`REWARD_KEY`
`LOSS_KEY`
`GRAD_NORM_KEY`
`LEARNING_RATE_KEY`
`MEAN_GEN_TOKENS_PER_SAMPLE_KEY`
`TRAINING_TEED_METRICS`

API#

nemo_rl.algorithms.metric_utils.SETUP_TIMING_PREFIX#: ‘timing/setup’

nemo_rl.algorithms.metric_utils.REWARD_KEY#: ‘reward’

nemo_rl.algorithms.metric_utils.LOSS_KEY#: ‘loss’

nemo_rl.algorithms.metric_utils.GRAD_NORM_KEY#: ‘grad_norm’

nemo_rl.algorithms.metric_utils.LEARNING_RATE_KEY#: ‘lr’

nemo_rl.algorithms.metric_utils.MEAN_GEN_TOKENS_PER_SAMPLE_KEY#: ‘mean_gen_tokens_per_sample’

nemo_rl.algorithms.metric_utils.TRAINING_TEED_METRICS#: ()

class nemo_rl.algorithms.metric_utils.SetupTimingMetrics#

Driver-side per-phase timings collected during setup.

generation_init_time_s: Optional[float]#: None

generation_init_reserve_time_s: Optional[float]#: None

generation_init_load_time_s: Optional[float]#: None

policy_init_time_s: Optional[float]#: None

value_init_time_s: Optional[float]#: None

nemo_gym_init_time_s: Optional[float]#: None

collective_init_time_s: Optional[float]#: None

weight_sync_time_s: Optional[float]#: None

parallel_wall_time_s: Optional[float]#: None

parallel_init_enabled: Optional[float]#: None

teacher_reservation_time_s: Optional[float]#: None

teacher_model_init_time_s: Optional[float]#: None

teacher_init_time_s: Optional[float]#: None

vllm_checkpoint_engine_init_time_s: Optional[float]#: None

total_setup_time_s: Optional[float]#: None

worker_setup_time_s: Optional[float]#: None

other_setup_time_s: Optional[float]#: None

extras: dict[str, float]#: ‘field(…)’

_DURATION_SUFFIXES#: (‘_time_s’, ‘_s’)

to_metrics_dict() → dict[str, Any]#: Serialize for Logger.log_metrics; drops unset (None) fields.

classmethod phase_name(key: str) → Optional[str]#

Phase a serialized key names, or None when it is not a duration.

Both suffixes are stripped so the phase reads as generation_init rather than generation_init_time_s. Requiring one of them is also the filter: it excludes non-durations that ride the same dict, such as parallel_init_enabled.

classmethod phase_seconds(metrics: dict[str, Any]) → dict[str, float]#

Extract {phase: seconds} from a serialized timing dict.

Reads a plain dict rather than an instance because extras and ppo.py’s hand-built dict of the same shape have to go through it too.

nemo_rl.algorithms.metric_utils.print_setup_timing_summary( metrics: nemo_rl.algorithms.metric_utils.SetupTimingMetrics, ) → None#