rlightning.utils.profiler.profiler

rlightning.utils.profiler.profiler.log_gpu_memory_usage(head: str, level=20, rank: int = 0)[source]

Log GPU memory usage information.

Parameters:

  • head (str) – A descriptive header for the memory usage log message.

  • logger (logging.Logger, optional) – Logger instance to use for logging. If None, prints to stdout.

  • level – Logging level to use. Defaults to logging.DEBUG.

  • rank (int) – The rank of the process to log memory for. Defaults to 0.

rlightning.utils.profiler.profiler.record_timing(name: str, duration_s: float, timing_raw: dict[str, _TimingStat], level: str = 'info')[source]

Record a single timing value into timing statistics storage.

Parameters:

  • name (str) – Name identifier for the timing measurement.

  • duration_s (float) – Duration in seconds to record. Must be a numeric value.

  • timing_raw (TimingRaw) – Dictionary to store timing statistics with the same structure as used by timer.

  • level (str) – Logging level for the metric, either “info” or “debug”. Defaults to “info”.

rlightning.utils.profiler.profiler.timer(name: str, timing_raw: dict[str, _TimingStat], level: str = 'info', enable: bool = True)[source]

Context manager for timing code blocks and recording timing statistics.

Parameters:

  • name (str) – Name identifier for the timing measurement.

  • timing_raw (TimingRaw) – Dictionary to store timing statistics with keys as names and values as dicts containing count, total, and avg.

  • level (str) – Logging level for the metric, either “info” or “debug”. Defaults to “info”.

  • enable (bool) – When False, the context manager does nothing. Defaults to True.

rlightning.utils.profiler.profiler.timer_wrap(name: str | None = None, level: str = 'info', log_to_metric: bool = False, enable: bool = True)[source]

Decorator for class methods to automatically time execution and record statistics.

Parameters:

  • name (str | None) – Name identifier for the timing measurement. If None, uses the function name. Defaults to None.

  • level (str) – Logging level for the metric, either “info” or “debug”. Defaults to “info”.

  • log_to_metric (bool) – When True, logging as metrics, otherwise logging to console.

  • enable (bool) – When False, the context manager does nothing. Defaults to True.

Returns:

The decorator function that wraps the original method.

Return type:

function