Add AggregatedTimer implementation (#1013)

Summary: Pull Request resolved: #1013

Reviewed By: JKSenthil

Differential Revision: D77345641

fbshipit-source-id: 72c017887e0a5c0cf8627442c3d2d6e2a38a1413

Author: diego-urgell

tests/utils/test_timer.py

@@ -19,6 +19,7 @@
 from torchtnt.utils.distributed import spawn_multi_process
 from torchtnt.utils.test_utils import skip_if_not_distributed
 from torchtnt.utils.timer import (
+    AggregatedTimer,
     BoundedTimer,
     FullSyncPeriodicTimer,
     get_durations_histogram,
@@ -235,6 +236,69 @@ def test_get_recorded_durations_table(self) -> None:
             == "\n| Name   |   p50 |   p90 |\n|:-------|------:|------:|\n| op     |     1 |     2 |"
         )
 
+    def test_aggregated_timer_statistics(self) -> None:
+        """Test that AggregatedTimer maintains correct statistics"""
+        timer = AggregatedTimer()
+
+        # Time an action multiple times with known durations
+        sleep_times = [1.0, 1.5, 2, 2.5, 1.0]
+        for sleep_time in sleep_times:
+            with timer.time("test_action"):
+                time.sleep(sleep_time)
+
+        with timer.time("fast_action"):
+            time.sleep(0.05)
+
+        # Generate report to calculate mean_duration
+        report = timer._make_report()
+        # test_action should be first
+        test_action_report = report.timed_action_stats[0]
+        fast_action_report = report.timed_action_stats[1]
+
+        # Verify for test_action
+        self.assertEqual(test_action_report.num_calls, 5)
+        expected_total = sum(sleep_times)
+        self.assert_within_tolerance(
+            test_action_report.total_duration, expected_total, 20
+        )
+        self.assert_within_tolerance(
+            test_action_report.mean_duration, expected_total / 5
+        )
+
+        # Verify for fast_action
+        self.assertEqual(fast_action_report.num_calls, 1)
+        self.assert_within_tolerance(fast_action_report.total_duration, 0.05)
+        self.assertGreater(
+            test_action_report.mean_duration, fast_action_report.mean_duration
+        )
+
+    def test_aggregated_timer_multiple_actions(self) -> None:
+        timer = AggregatedTimer()
+
+        actions = ["dataloader", "forward", "backward", "optimizer"]
+        call_counts = [10, 10, 10, 5]
+
+        for action, count in zip(actions, call_counts):
+            for _ in range(count):
+                with timer.time(action):
+                    time.sleep(0.001)
+
+        # Check that all actions are recorded
+        self.assertEqual(len(timer._aggregate_stats), 4)
+
+        for action, expected_count in zip(actions, call_counts):
+            self.assertEqual(timer._aggregate_stats[action].num_calls, expected_count)
+
+        # Check report
+        report = timer._make_report()
+        self.assertEqual(len(report.timed_action_stats), 4)
+        self.assertEqual(report.total_calls, sum(call_counts))
+
+        total_percentage = sum(
+            stats.percentage_of_total_time for stats in report.timed_action_stats
+        )
+        self.assert_within_tolerance(total_percentage, 100.0, 1)
+
 
 class FullSyncPeriodicTimerTest(unittest.TestCase):
     @classmethod

torchtnt/utils/timer.py

@@ -208,6 +208,75 @@ def _make_report(self: TimerProtocol) -> TimerReport:
         )
 
 
+class AggregatedTimer(Timer):
+    """
+    A Timer class which implements TimerProtocol and stores aggregated timing stats. Instead of
+    storing the recorded durations for each action, it accumulates running metrics and computes
+    final stats at the end when generating the report. This is useful for cases where the number
+    of samples is too large to store in memory.
+    """
+
+    def __init__(
+        self,
+        cuda_sync: Optional[bool] = None,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(cuda_sync=cuda_sync, verbose=verbose)
+        self._aggregate_stats: Dict[str, TimedActionStats] = defaultdict(
+            lambda: TimedActionStats(action_name="")  # Filled in on time() method
+        )
+
+    @contextmanager
+    def time(
+        self,
+        action_name: str,
+    ) -> Generator[None, None, None]:
+        # Run base class context manager first
+        with super().time(action_name):
+            yield
+
+        # Update aggregate stats
+        latest_duration: float = self.recorded_durations[action_name][-1]
+        self._aggregate_stats[action_name].action_name = action_name
+        self._aggregate_stats[action_name].num_calls += 1
+        self._aggregate_stats[action_name].total_duration += latest_duration
+
+        # Reset recorded durations to avoid storing data
+        self.recorded_durations.clear()
+
+    def _make_report(self) -> TimerReport:
+        """
+        Creates the report but considering that the data is aggregated in the correct structure.
+        """
+        total_time = 0.0
+        total_calls = 0
+
+        # Calculate total time and calls across all actions
+        for stats in self._aggregate_stats.values():
+            total_time += stats.total_duration
+            total_calls += stats.num_calls
+
+        # Build report data
+        action_stats: List[TimedActionStats] = []
+        for stats in self._aggregate_stats.values():
+            stats.mean_duration = (
+                stats.total_duration / stats.num_calls if stats.num_calls > 0 else 0.0
+            )
+            stats.percentage_of_total_time = (
+                100.0 * stats.total_duration / total_time if total_time > 0 else 0.0
+            )
+            action_stats.append(stats)
+
+        # Sort by percentage (descending)
+        action_stats.sort(key=lambda x: x.percentage_of_total_time, reverse=True)
+
+        return TimerReport(
+            timed_action_stats=action_stats,
+            total_calls=total_calls,
+            total_duration=total_time,
+        )
+
+
 class BoundedTimer(Timer):
     """
     A Timer class which implements TimerProtocol and stores timings in a dictionary `recorded_durations`.