nemo_rl.utils.timer#

Module Contents#

Classes#

Timer

A utility for timing code execution.

ThreadSafeTimer

Thread-safe extension of Timer for use in multi-threaded contexts.

TimeoutChecker

Functions#

convert_to_seconds

Converts a time string in the format ‘DD:HH:MM:SS’ to total seconds.

Data#

logger

API#

nemo_rl.utils.timer.logger#

‘getLogger(…)’

class nemo_rl.utils.timer.Timer(context: Optional[dict[str, object]] = None)#

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

Initialize the timer.

Parameters:

context – Arbitrary key-value pairs identifying this timer instance. Included in DEBUG log messages emitted on every start/stop/record/mark call. Typical keys: rank, worker, node, job_id. Example: {“rank”: 3, “worker”: “collector”, “node”: “gpu-05”}

_REDUCTION_FUNCTIONS: dict[str, Callable[[Sequence[float]], float]]#

None

_fmt(label: str, event: str) → str#

Build a log message string, prepending context prefix and appending UTC timestamp.

start(label: str, should_log: bool = True) → None#

Start timing for the given label.

stop(label: str, should_log: bool = True) → 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

record(label: str, elapsed: float) → None#

Append a pre-measured duration without start/stop.

Useful when the caller has already measured the elapsed time (e.g., across a try/except boundary) and wants to record it directly.

Parameters:

  • label – The timing label to record under

  • elapsed – The elapsed time in seconds

mark(label: str, metadata: Optional[dict] = None) → float#

Record a point-in-time event at the current Unix epoch.

Unlike start/stop/record which measure durations, this captures a standalone timestamp for events like failures, state transitions, or any moment worth noting on a timeline.

Uses time.time() (not perf_counter) so timestamps are correlatable across processes and Ray actors.

Parameters:

  • label – The event label (e.g., “vllm/worker_crashed”)

  • metadata – Optional context dict (e.g., {“worker_id”: 3, “error”: “OOM”})

Returns:

The recorded timestamp (Unix epoch seconds).

get_markers(
label: Optional[str] = None,
) → dict[str, list[tuple[float, Optional[dict]]]]#

Get recorded markers, optionally filtered by label.

Parameters:

label – If provided, return markers for this label only. If None, return all markers.

Returns:

Dict mapping labels to lists of (timestamp, metadata) tuples.

time(
label: str,
should_log: bool = True,
) → Generator[None, None, None]#

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, float | 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 and markers for the specified label or all labels.

Parameters:

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

class nemo_rl.utils.timer.ThreadSafeTimer(context: Optional[dict[str, object]] = None)#

Bases: nemo_rl.utils.timer.Timer

Thread-safe extension of Timer for use in multi-threaded contexts.

Wraps all mutating and reading operations with a lock so that concurrent threads can safely record timings to the same instance.

Initialization

Initialize the timer.

Parameters:

context – Arbitrary key-value pairs identifying this timer instance. Included in DEBUG log messages emitted on every start/stop/record/mark call. Typical keys: rank, worker, node, job_id. Example: {“rank”: 3, “worker”: “collector”, “node”: “gpu-05”}

start(label: str, should_log: bool = True) → None#
stop(label: str, should_log: bool = True) → float#
record(label: str, elapsed: float) → None#
mark(label: str, metadata: Optional[dict] = None) → float#
get_markers(
label: Optional[str] = None,
) → dict[str, list[tuple[float, Optional[dict]]]]#
time(
label: str,
should_log: bool = True,
) → Generator[None, None, None]#
get_elapsed(label: str) → list[float]#
get_latest_elapsed(label: str) → float#
reduce(label: str, operation: str = 'mean') → float#
get_timing_metrics(
reduction_op: Union[str, dict[str, str]] = 'mean',
) → dict[str, float | list[float]]#
reset(label: Optional[str] = None) → None#
nemo_rl.utils.timer.convert_to_seconds(time_string: str) → int#

Converts a time string in the format ‘DD:HH:MM:SS’ to total seconds.

Parameters:

time_string (str) – Time duration string, e.g., ‘00:03:45:00’.

Returns:

Total time in seconds.

Return type:

int

class nemo_rl.utils.timer.TimeoutChecker(
timeout: Optional[str] = '00:03:45:00',
fit_last_save_time: bool = False,
)#

Initialization

Initializes the TimeoutChecker.

Parameters:

  • timeout (str or None) – Timeout in format ‘DD:HH:MM:SS’. If None, timeout is considered infinite.

  • fit_last_save_time (bool) – If True, considers average iteration time when checking timeout.

check_save()#
start_iterations()#
mark_iteration()#