From 61aaae7e534719e0d285036b937195defc87c6b1 Mon Sep 17 00:00:00 2001 From: Diego Urgell Date: Thu, 26 Jun 2025 13:59:45 -0700 Subject: [PATCH 1/2] Refactor _make_report methods with better interface (#1011) Summary: Pull Request resolved: https://github.com/pytorch/tnt/pull/1011 Reviewed By: JKSenthil Differential Revision: D77345640 --- torchtnt/utils/timer.py | 100 +++++++++++++++++++++++----------------- 1 file changed, 57 insertions(+), 43 deletions(-) diff --git a/torchtnt/utils/timer.py b/torchtnt/utils/timer.py index b25a9b6ac9..b9a25591a5 100644 --- a/torchtnt/utils/timer.py +++ b/torchtnt/utils/timer.py @@ -12,6 +12,8 @@ import os from collections import defaultdict from contextlib import contextmanager +from dataclasses import dataclass +from functools import total_ordering from time import perf_counter from typing import ( Any, @@ -22,7 +24,6 @@ Protocol, runtime_checkable, Sequence, - Tuple, ) import numpy as np @@ -33,10 +34,8 @@ from torch.distributed.distributed_c10d import Work from torchtnt.utils.distributed import PGWrapper -logger: logging.Logger = logging.getLogger(__name__) -_TABLE_ROW = Tuple[str, float, int, float, float] -_TABLE_DATA = List[_TABLE_ROW] +logger: logging.Logger = logging.getLogger(__name__) @contextmanager @@ -69,6 +68,28 @@ def log_elapsed_time( logger.info(f"{action_name} took {interval_time} seconds") +@total_ordering +@dataclass +class TimedActionStats: + """Dataclass for storing timed action stats. These can be consumed by report generation methods, so metrics should be aggregated.""" + + action_name: str + mean_duration: float = 0.0 + num_calls: int = 0 + total_duration: float = 0.0 + percentage_of_total_time: float = 0.0 + + def __le__(self, other: "TimedActionStats") -> bool: + return self.percentage_of_total_time <= other.percentage_of_total_time + + +@dataclass +class TimerReport: + timed_action_stats: List[TimedActionStats] + total_calls: int + total_duration: float + + @runtime_checkable class TimerProtocol(Protocol): """ @@ -194,31 +215,30 @@ def _apply_bounds(self, action_name: str) -> None: ) -def _get_total_time(timer: TimerProtocol) -> float: +def _make_report(self: TimerProtocol) -> TimerReport: total_time = 0.0 - for _, durations in timer.recorded_durations.items(): + for _, durations in self.recorded_durations.items(): array_value = np.array(durations) array_sum = np.sum(array_value) total_time += array_sum - return total_time - - -def _make_report(timer: TimerProtocol) -> Tuple[_TABLE_DATA, float, float]: - total_time = _get_total_time(timer) - report = [ - ( - a, - np.mean(d), - len(d), - np.sum(d), - 100.0 * np.sum(d) / total_time, + action_stats = [ + TimedActionStats( + action_name=a, + mean_duration=np.mean(d), + num_calls=len(d), + total_duration=np.sum(d), + percentage_of_total_time=100.0 * np.sum(d) / total_time, ) - for a, d in timer.recorded_durations.items() + for a, d in self.recorded_durations.items() ] - report.sort(key=lambda x: x[4], reverse=True) - total_calls = sum(x[2] for x in report) - return report, total_calls, total_time + action_stats.sort(reverse=True) + total_calls = sum(x.num_calls for x in action_stats) + return TimerReport( + timed_action_stats=action_stats, + total_calls=total_calls, + total_duration=total_time, + ) def get_timer_summary(timer: TimerProtocol) -> str: @@ -231,13 +251,16 @@ def get_timer_summary(timer: TimerProtocol) -> str: ValueError If the input Timer has no recorded actions """ + report: TimerReport = _make_report(timer) + sep: str = os.linesep output_string = f"Timer Report{sep}" - if len(timer.recorded_durations) == 0: + # Handle empty timer case + if not report.timed_action_stats: return output_string - max_key = max(len(k) for k in timer.recorded_durations.keys()) + max_key = max(len(a.action_name) for a in report.timed_action_stats) # pyre-fixme[53]: Captured variable `max_key` is not annotated. def log_row(action: str, mean: str, num_calls: str, total: str, per: str) -> str: @@ -252,32 +275,23 @@ def log_row(action: str, mean: str, num_calls: str, total: str, per: str) -> str "Total time (s)", "Percentage %", ) + output_string_len = len(header_string.expandtabs()) - 1 sep_lines = f"{sep}{'-' * output_string_len}" output_string += sep_lines + header_string + sep_lines - report: _TABLE_DATA - ( - report, - total_calls, - total_duration, - ) = _make_report(timer) + output_string += log_row( - "Total", "-", f"{total_calls:}", f"{total_duration:.5}", "100 %" + "Total", "-", f"{report.total_calls:}", f"{report.total_duration:.5}", "100 %" ) output_string += sep_lines - for ( - action, - mean_duration, - num_calls, - total_duration, - duration_per, - ) in report: + + for action in report.timed_action_stats: output_string += log_row( - action, - f"{mean_duration:.5}", - f"{num_calls}", - f"{total_duration:.5}", - f"{duration_per:.5}", + action.action_name, + f"{action.mean_duration:.5}", + f"{action.num_calls}", + f"{action.total_duration:.5}", + f"{action.percentage_of_total_time:.5}", ) output_string += sep_lines From a6ea72d24b5c583b1052837152921add31e706ee Mon Sep 17 00:00:00 2001 From: Diego Urgell Date: Thu, 26 Jun 2025 13:59:45 -0700 Subject: [PATCH 2/2] Add report generation methods to TimerProtocol interface (#1012) Summary: Pull Request resolved: https://github.com/pytorch/tnt/pull/1012 Reviewed By: JKSenthil Differential Revision: D77345642 --- torchtnt/utils/timer.py | 60 ++++++++++++++++++++++------------------- 1 file changed, 32 insertions(+), 28 deletions(-) diff --git a/torchtnt/utils/timer.py b/torchtnt/utils/timer.py index b9a25591a5..ac3226ce54 100644 --- a/torchtnt/utils/timer.py +++ b/torchtnt/utils/timer.py @@ -114,6 +114,12 @@ def reset(self) -> None: """ ... + def _make_report(self) -> TimerReport: + """ + Creates a report of timing data. + """ + ... + class Timer(TimerProtocol): def __init__( @@ -176,6 +182,31 @@ def reset(self) -> None: """ self.recorded_durations = defaultdict(list) + def _make_report(self: TimerProtocol) -> TimerReport: + total_time = 0.0 + for _, durations in self.recorded_durations.items(): + array_value = np.array(durations) + array_sum = np.sum(array_value) + total_time += array_sum + + action_stats = [ + TimedActionStats( + action_name=a, + mean_duration=np.mean(d), + num_calls=len(d), + total_duration=np.sum(d), + percentage_of_total_time=100.0 * np.sum(d) / total_time, + ) + for a, d in self.recorded_durations.items() + ] + action_stats.sort(reverse=True) + total_calls = sum(x.num_calls for x in action_stats) + return TimerReport( + timed_action_stats=action_stats, + total_calls=total_calls, + total_duration=total_time, + ) + class BoundedTimer(Timer): """ @@ -215,32 +246,6 @@ def _apply_bounds(self, action_name: str) -> None: ) -def _make_report(self: TimerProtocol) -> TimerReport: - total_time = 0.0 - for _, durations in self.recorded_durations.items(): - array_value = np.array(durations) - array_sum = np.sum(array_value) - total_time += array_sum - - action_stats = [ - TimedActionStats( - action_name=a, - mean_duration=np.mean(d), - num_calls=len(d), - total_duration=np.sum(d), - percentage_of_total_time=100.0 * np.sum(d) / total_time, - ) - for a, d in self.recorded_durations.items() - ] - action_stats.sort(reverse=True) - total_calls = sum(x.num_calls for x in action_stats) - return TimerReport( - timed_action_stats=action_stats, - total_calls=total_calls, - total_duration=total_time, - ) - - def get_timer_summary(timer: TimerProtocol) -> str: """Given a timer, generate a summary of all the recorded actions. @@ -251,7 +256,7 @@ def get_timer_summary(timer: TimerProtocol) -> str: ValueError If the input Timer has no recorded actions """ - report: TimerReport = _make_report(timer) + report: TimerReport = timer._make_report() sep: str = os.linesep output_string = f"Timer Report{sep}" @@ -279,7 +284,6 @@ def log_row(action: str, mean: str, num_calls: str, total: str, per: str) -> str output_string_len = len(header_string.expandtabs()) - 1 sep_lines = f"{sep}{'-' * output_string_len}" output_string += sep_lines + header_string + sep_lines - output_string += log_row( "Total", "-", f"{report.total_calls:}", f"{report.total_duration:.5}", "100 %" )