add progress to error messages (#779)

Summary: Pull Request resolved: #779

Reviewed By: JKSenthil

Differential Revision: D55891686

fbshipit-source-id: 15741a2d7cea147169dc0a9b4ceb4cafe44766db

Author: galrotem

tests/framework/test_evaluate.py

@@ -205,6 +205,16 @@ def test_evaluate_timing(self) -> None:
         )
         self.assertIn("evaluate.next(data_iter)", timer.recorded_durations.keys())
 
+    def test_error_message(self) -> None:
+        with self.assertRaises(ValueError), self.assertLogs(level="INFO") as log:
+            evaluate(EvalUnitWithError(), [1, 2, 3, 4])
+
+        self.assertIn(
+            "INFO:torchtnt.framework.evaluate:Exception during evaluate after the following progress: "
+            "completed epochs: 0, completed steps: 1, completed steps in current epoch: 1.:\nfoo",
+            log.output,
+        )
+
 
 class StopEvalUnit(EvalUnit[Tuple[torch.Tensor, torch.Tensor]]):
     def __init__(self, input_dim: int, steps_before_stopping: int) -> None:
@@ -235,4 +245,7 @@ def eval_step(
         return loss, outputs
 
 
-Batch = Tuple[torch.Tensor, torch.Tensor]
+class EvalUnitWithError(EvalUnit[int]):
+    def eval_step(self, state: State, data: int) -> None:
+        if self.eval_progress.num_steps_completed == 1:
+            raise ValueError("foo")

tests/framework/test_fit.py

@@ -328,3 +328,30 @@ def test_fit_timing(self) -> None:
         )
         self.assertIn("train.next(data_iter)", timer.recorded_durations.keys())
         self.assertIn("evaluate.next(data_iter)", timer.recorded_durations.keys())
+
+    def test_error_message(self) -> None:
+        self.maxDiff = None
+        with self.assertRaises(ValueError), self.assertLogs(level="INFO") as log:
+            fit(
+                UnitWithError(),
+                train_dataloader=[1, 2, 3, 4],
+                eval_dataloader=[1, 2, 3, 4],
+                max_steps=10,
+                evaluate_every_n_epochs=1,
+            )
+
+        self.assertIn(
+            "INFO:torchtnt.framework.fit:Exception during fit after the following progress: train "
+            "progress: completed epochs: 1, completed steps: 4, completed steps in current epoch: 0. "
+            "eval progress: completed epochs: 0, completed steps: 2, completed steps in current epoch: 2.:\nfoo",
+            log.output,
+        )
+
+
+class UnitWithError(TrainUnit[int], EvalUnit[int]):
+    def train_step(self, state: State, data: int) -> None:
+        pass
+
+    def eval_step(self, state: State, data: int) -> None:
+        if self.eval_progress.num_steps_completed == 2:
+            raise ValueError("foo")

tests/framework/test_predict.py

@@ -213,6 +213,16 @@ def test_predict_timing(self) -> None:
         )
         self.assertIn("predict.next(data_iter)", timer.recorded_durations.keys())
 
+    def test_error_message(self) -> None:
+        with self.assertRaises(ValueError), self.assertLogs(level="INFO") as log:
+            predict(PredictUnitWithError(), [1, 2, 3, 4])
+
+        self.assertIn(
+            "INFO:torchtnt.framework.predict:Exception during predict after the following progress: "
+            "completed epochs: 0, completed steps: 3, completed steps in current epoch: 3.:\nfoo",
+            log.output,
+        )
+
 
 Batch = Tuple[torch.Tensor, torch.Tensor]
 
@@ -238,3 +248,9 @@ def predict_step(self, state: State, data: Batch) -> torch.Tensor:
 
         self.steps_processed += 1
         return outputs
+
+
+class PredictUnitWithError(PredictUnit[int]):
+    def predict_step(self, state: State, data: int) -> None:
+        if self.predict_progress.num_steps_completed == 3:
+            raise ValueError("foo")

tests/framework/test_train.py

@@ -263,6 +263,16 @@ def test_train_timing(self) -> None:
         )
         self.assertIn("train.next(data_iter)", timer.recorded_durations.keys())
 
+    def test_error_message(self) -> None:
+        with self.assertRaises(ValueError), self.assertLogs(level="INFO") as log:
+            train(TrainUnitWithError(), [1, 2, 3, 4], max_steps=10)
+
+        self.assertIn(
+            "INFO:torchtnt.framework.train:Exception during train after the following progress: "
+            "completed epochs: 0, completed steps: 2, completed steps in current epoch: 2.:\nfoo",
+            log.output,
+        )
+
 
 Batch = Tuple[torch.Tensor, torch.Tensor]
 
@@ -300,4 +310,10 @@ def train_step(
         return loss, outputs
 
 
+class TrainUnitWithError(TrainUnit[Batch]):
+    def train_step(self, state: State, data: Batch) -> None:
+        if self.train_progress.num_steps_completed == 2:
+            raise ValueError("foo")
+
+
 Batch = Tuple[torch.Tensor, torch.Tensor]

torchtnt/framework/evaluate.py

@@ -95,7 +95,9 @@ def evaluate(
             logger.info(get_timer_summary(state.timer))
     except Exception as e:
         # TODO: log for diagnostics
-        logger.info(e)
+        logger.info(
+            f"Exception during evaluate after the following progress: {eval_unit.eval_progress.get_progress_string()}:\n{e}"
+        )
         eval_unit.on_exception(state, e)
         callback_handler.on_exception(state, eval_unit, e)
         raise e

torchtnt/framework/fit.py

@@ -136,7 +136,9 @@ def fit(
             logger.info(get_timer_summary(state.timer))
     except Exception as e:
         # TODO: log for diagnostics
-        logger.info(f"Exception during fit:\n{e}")
+        logger.info(
+            f"Exception during fit after the following progress: train progress: {unit.train_progress.get_progress_string()} eval progress: {unit.eval_progress.get_progress_string()}:\n{e}"
+        )
         unit.on_exception(state, e)
         callback_handler.on_exception(state, unit, e)
         raise e

torchtnt/framework/predict.py

@@ -95,7 +95,9 @@ def predict(
             logger.info(get_timer_summary(state.timer))
     except Exception as e:
         # TODO: log for diagnostics
-        logger.info(f"Exception during predict:\n{e}")
+        logger.info(
+            f"Exception during predict after the following progress: {predict_unit.predict_progress.get_progress_string()}:\n{e}"
+        )
         predict_unit.on_exception(state, e)
         callback_handler.on_exception(state, predict_unit, e)
         raise e

torchtnt/framework/train.py

@@ -105,7 +105,9 @@ def train(
             logger.info(get_timer_summary(state.timer))
     except Exception as e:
         # TODO: log for diagnostics
-        logger.info(f"Exception during train:\n{e}")
+        logger.info(
+            f"Exception during train after the following progress: {train_unit.train_progress.get_progress_string()}:\n{e}"
+        )
         train_unit.on_exception(state, e)
         callback_handler.on_exception(state, train_unit, e)
         raise e

torchtnt/utils/progress.py

@@ -62,6 +62,9 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         self._num_steps_completed = state_dict["num_steps_completed"]
         self._num_steps_completed_in_epoch = state_dict["num_steps_completed_in_epoch"]
 
+    def get_progress_string(self) -> str:
+        return f"completed epochs: {self.num_epochs_completed}, completed steps: {self.num_steps_completed}, completed steps in current epoch: {self.num_steps_completed_in_epoch}."
+
 
 def estimated_steps_in_epoch(
     dataloader: Iterable[object],