Add report generation methods to TimerProtocol interface

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):
     """
@@ -93,6 +114,12 @@ def reset(self) -> None:
         """
         ...
 
+    def _make_report(self) -> TimerReport:
+        """
+        Creates a report of timing data.
+        """
+        ...
+
 
 class Timer(TimerProtocol):
     def __init__(
@@ -155,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):
     """
@@ -194,33 +246,6 @@ def _apply_bounds(self, action_name: str) -> None:
             )
 
 
-def _get_total_time(timer: TimerProtocol) -> float:
-    total_time = 0.0
-    for _, durations in timer.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,
-        )
-        for a, d in timer.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
-
-
 def get_timer_summary(timer: TimerProtocol) -> str:
     """Given a timer, generate a summary of all the recorded actions.
 
@@ -231,13 +256,16 @@ def get_timer_summary(timer: TimerProtocol) -> str:
         ValueError
             If the input Timer has no recorded actions
     """
+    report: TimerReport = timer._make_report()
+
     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 +280,22 @@ 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