Fix ThroughputLogger error on intra-epoch restore (#829)

Summary: Pull Request resolved: #829

Reviewed By: galrotem, JKSenthil

Differential Revision: D57183157

fbshipit-source-id: 839dfb149a89d4c92c87f725a613b8e7104b499c

Author: diego-urgell

tests/framework/callbacks/test_throughput_logger.py

@@ -272,6 +272,10 @@ def test_log_for_epoch(self) -> None:
             throughput_logger.on_train_step_end(state, unit)
         self.assertEqual(throughput_logger._steps_in_epoch[ActivePhase.TRAIN], 1)
 
+        # Make sure we don't log or fail if the _epoch_start_times dict is not initialized
+        throughput_logger._log_for_epoch(state, epoch_logging_for=15)
+        logger.log.assert_not_called()
+
         with patch("time.perf_counter", return_value=0.5):
             throughput_logger.on_train_epoch_start(state, MagicMock(spec=TrainUnit))
         self.assertEqual(throughput_logger._epoch_start_times[ActivePhase.TRAIN], 0.5)

torchtnt/framework/callbacks/throughput_logger.py

@@ -7,6 +7,7 @@
 # pyre-strict
 
 
+import logging
 import time
 from collections import defaultdict
 from typing import Dict, Mapping
@@ -30,6 +31,8 @@
     ActivePhase.PREDICT: "Predict",
 }
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 class ThroughputLogger(Callback):
     """
@@ -192,6 +195,18 @@ def _log_for_epoch(
         *,
         epoch_logging_for: int,
     ) -> None:
+
+        # Avoid key errors if active phase is not set. This may happen if we restore
+        # from an intra-epoch checkpoint.
+        if (
+            state.active_phase not in self._epoch_start_times
+            or state.active_phase not in self._steps_in_epoch
+        ):
+            logger.warning(
+                f"Missing troughput data for epoch {epoch_logging_for}, phase {state.active_phase}. Ommiting troughput logging."
+            )
+            return
+
         time_since_epoch_start = (
             time.perf_counter() - self._epoch_start_times[state.active_phase]
         )