nemo_rl.utils.timer

Module Contents

Classes

Timer

A utility for timing code execution.

API

class nemo_rl.utils.timer.Timer

A utility for timing code execution.

Supports two usage patterns:

1. Explicit start/stop: timer.start(“label”), timer.stop(“label”)

2. Context manager: with timer.time(“label”): …

The timer keeps track of multiple timing measurements for each label, and supports different reductions on these measurements (mean, median, min, max, std dev).

Example usage:

timer = Timer()

# Method 1: start/stop
timer.start("load_data")
data = load_data()
timer.stop("load_data")

# Method 2: context manager
with timer.time("model_forward"):
    model_outputs = model(inputs)

# Multiple timing measurements for the same operation
for batch in dataloader:
    with timer.time("model_forward_multiple"):
        outputs = model(batch)

# Get all times for one label
model_forward_times = timer.get_elapsed("model_forward_multiple")

# Get reductions for one label
mean_forward_time = timer.reduce("model_forward_multiple")
max_forward_time = timer.reduce("model_forward_multiple", "max")

Initialization

_REDUCTION_FUNCTIONS

None

start(label: str) → None

Start timing for the given label.

stop(label: str) → float

Stop timing for the given label and return the elapsed time.

Parameters:

label – The label to stop timing for

Returns:

The elapsed time in seconds

Raises:

ValueError – If the timer for the given label is not running

time(label: str)

Context manager for timing a block of code.

Parameters:

label – The label to use for this timing

Yields:

None

get_elapsed(label: str) → List[float]

Get all elapsed time measurements for a specific label.

Parameters:

label – The timing label to get elapsed times for

Returns:

A list of all elapsed time measurements in seconds

Raises:

KeyError – If the label doesn’t exist

get_latest_elapsed(label: str) → float

Get the most recent elapsed time measurement for a specific label.

Parameters:

label – The timing label to get the latest elapsed time for

Returns:

The most recent elapsed time measurement in seconds

Raises:

• KeyError – If the label doesn’t exist

• IndexError – If the label exists but has no measurements

reduce(label: str, operation: str = 'mean') → float

Apply a reduction function to timing measurements for the specified label.

Parameters:

• label – The timing label to get reduction for

• operation –

  The type of reduction to apply. Valid options are:

  • ”mean”: Average time (default)

  • ”median”: Median time

  • ”min”: Minimum time

  • ”max”: Maximum time

  • ”std”: Standard deviation

  • ”sum”: Total time

  • ”count”: Number of measurements

Returns:

A single float with the reduction result

Raises:

• KeyError – If the label doesn’t exist

• ValueError – If an invalid operation is provided

get_timing_metrics(

reduction_op: Union[str, Dict[str, str]] = 'mean',

) → Dict[str, List[float]]

Get all timing measurements with optional reduction.

Parameters:

reduction_op – Either a string specifying a reduction operation to apply to all labels, or a dictionary mapping specific labels to reduction operations. Valid reduction operations are: “mean”, “median”, “min”, “max”, “std”, “sum”, “count”. If a label is not in the dictionary, no reduction is applied and all measurements are returned.

Returns:

• A list of all timing measurements for that label (if no reduction specified)

• A single float with the reduction result (if reduction specified)

Return type:

A dictionary mapping labels to either

Raises:

ValueError – If an invalid reduction operation is provided

reset(label: Optional[str] = None) → None

Reset timings for the specified label or all labels.

Parameters:

label – Optional label to reset. If None, resets all timers.