Unify current_fx_name and current_hook_fx_name [2/n] (#7594)

* Minor loggger connector cleanup [1/n]

* Missing line

* Address comments

* Rely on validator

* Unify `current_fx_name` and `current_hook_fx_name`

* Fix test

Author: carmocca

pytorch_lightning/core/lightning.py

@@ -104,9 +104,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._example_input_array = None
         self._datamodule = None
         self._results: Optional[Result] = None
-        self._current_fx_name: str = ''
+        self._current_fx_name: Optional[str] = None
         self._running_manual_backward: bool = False
-        self._current_hook_fx_name: Optional[str] = None
         self._current_dataloader_idx: Optional[int] = None
         self._automatic_optimization: bool = True
         self._truncated_bptt_steps: int = 0
@@ -316,8 +315,9 @@ def log(
             on_step = self.__auto_choose_log_on_step(on_step)
             on_epoch = self.__auto_choose_log_on_epoch(on_epoch)
 
+            assert self._current_fx_name is not None
             self.trainer.logger_connector.check_logging_in_callbacks(
-                self._current_hook_fx_name, on_step=on_step, on_epoch=on_epoch
+                self._current_fx_name, on_step=on_step, on_epoch=on_epoch
             )
 
             # make sure user doesn't introduce logic for multi-dataloaders

pytorch_lightning/trainer/connectors/logger_connector/callback_hook_validator.py

@@ -18,8 +18,8 @@
 class CallbackHookNameValidator:
 
     @staticmethod
-    def check_logging_in_callbacks(current_hook_fx_name: str, on_step: bool, on_epoch: bool) -> None:
-        internal_func = getattr(CallbackHookNameValidator, f"_{current_hook_fx_name}_log", None)
+    def check_logging_in_callbacks(fx_name: str, on_step: bool, on_epoch: bool) -> None:
+        internal_func = getattr(CallbackHookNameValidator, f"_{fx_name}_log", None)
         if internal_func is None:
             return
 
@@ -28,16 +28,14 @@ def check_logging_in_callbacks(current_hook_fx_name: str, on_step: bool, on_epoc
         if current_callback_hook_auth_args is not None:
             m = "{} function supports only {} in {}. Provided {}"
             if on_step not in current_callback_hook_auth_args["on_step"]:
-                msg = m.format(current_hook_fx_name, "on_step", current_callback_hook_auth_args["on_step"], on_step)
+                msg = m.format(fx_name, "on_step", current_callback_hook_auth_args["on_step"], on_step)
                 raise MisconfigurationException(msg)
 
             if on_epoch not in current_callback_hook_auth_args["on_epoch"]:
-                msg = m.format(current_hook_fx_name, "on_epoch", current_callback_hook_auth_args["on_epoch"], on_epoch)
+                msg = m.format(fx_name, "on_epoch", current_callback_hook_auth_args["on_epoch"], on_epoch)
                 raise MisconfigurationException(msg)
         else:
-            raise MisconfigurationException(
-                f"{current_hook_fx_name} function doesn't support logging using self.log() yet."
-            )
+            raise MisconfigurationException(f"{fx_name} function doesn't support logging using self.log() yet.")
 
     @staticmethod
     def _on_before_accelerator_backend_setup_log():

pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py

@@ -228,8 +228,8 @@ class EpochResultStore:
 
     ..example::
 
-        model._current_fx_name = 'something'
         model._results = Result()
+        model._current_fx_name = 'something'
         model.log('a', ...)
         epoch_result_store.cache_result()
     """
@@ -250,7 +250,7 @@ def info(self):
         model_ref = self.trainer.lightning_module
         return {
             "batch_idx": self.trainer.train_loop.batch_idx,
-            "fx_name": model_ref._current_hook_fx_name or model_ref._current_fx_name,
+            "fx_name": model_ref._current_fx_name,
             "dataloader_idx": model_ref._current_dataloader_idx or 0,
             "opt_idx": self._opt_idx or 0,
             "split_idx": self._split_idx or 0,
@@ -266,8 +266,7 @@ def reset_model(self):
         """
         model_ref = self.trainer.lightning_module
         model_ref._results = Result()
-        model_ref._current_hook_fx_name = None
-        model_ref._current_fx_name = ''
+        model_ref._current_fx_name = None
 
     def cache_result(self) -> None:
         """
@@ -280,8 +279,7 @@ def cache_result(self) -> None:
             hook_result = model_ref._results
 
             if len(hook_result) == 1:
-                model_ref._current_hook_fx_name = None
-                model_ref._current_fx_name = ''
+                model_ref._current_fx_name = None
                 return
 
             info = self.info

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -95,10 +95,8 @@ def set_metrics(self, key: str, val: Dict) -> None:
     def reset(self) -> None:
         self.cached_results.reset()
 
-    def check_logging_in_callbacks(self, hook_fx_name: str, on_step: bool, on_epoch: bool) -> None:
-        self._callback_hook_validator.check_logging_in_callbacks(
-            current_hook_fx_name=hook_fx_name, on_step=on_step, on_epoch=on_epoch
-        )
+    def check_logging_in_callbacks(self, fx_name: str, on_step: bool, on_epoch: bool) -> None:
+        self._callback_hook_validator.check_logging_in_callbacks(fx_name=fx_name, on_step=on_step, on_epoch=on_epoch)
 
     def on_evaluation_batch_start(self, batch, dataloader_idx, num_dataloaders):
         model = self.trainer.lightning_module

pytorch_lightning/trainer/evaluation_loop.py

@@ -253,7 +253,7 @@ def on_evaluation_epoch_end(self) -> None:
         model_ref = self.trainer.lightning_module
         hook_name = "on_test_epoch_end" if self.trainer.testing else "on_validation_epoch_end"
 
-        self.trainer._reset_result_and_set_hook_fx_name(hook_name)
+        self.trainer._reset_result_and_set_fx_name(hook_name)
 
         with self.trainer.profiler.profile(hook_name):
 

pytorch_lightning/trainer/predict_loop.py

@@ -113,7 +113,7 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
 
         self.trainer.call_hook("on_predict_batch_start", batch, batch_idx, dataloader_idx)
 
-        model_ref._current_fx_name = "predict"
+        model_ref._current_fx_name = "predict_step"
         predictions = self.trainer.accelerator.predict_step(step_kwargs)
 
         if predictions is None:

pytorch_lightning/trainer/trainer.py

@@ -1186,19 +1186,18 @@ def _call_teardown_hook(self, model: LightningModule) -> None:
         self.teardown(stage=fn)
         model.teardown(stage=fn)
 
-        model._current_fx_name = ""
-        model._current_hook_fx_name = None
+        model._current_fx_name = None
         model._current_dataloader_idx = None
 
-    def _reset_result_and_set_hook_fx_name(self, hook_name: str) -> bool:
+    def _reset_result_and_set_fx_name(self, hook_name: str) -> bool:
         # on_before_zero_grad is called within training_step
         if "batch_start" in hook_name or hook_name in ("on_before_zero_grad", "on_after_backward"):
             return True
         model_ref = self.lightning_module
         if model_ref is not None:
             # used to track current hook name called
             model_ref._results = Result()
-            model_ref._current_hook_fx_name = hook_name
+            model_ref._current_fx_name = hook_name
         return False
 
     def _cache_logged_metrics(self):
@@ -1214,7 +1213,7 @@ def call_hook(self, hook_name: str, *args, **kwargs) -> Any:
         # TrainLoop._on_train_epoch_end_hook
 
         # set hook_name to model + reset Result obj
-        skip = self._reset_result_and_set_hook_fx_name(hook_name)
+        skip = self._reset_result_and_set_fx_name(hook_name)
 
         # always profile hooks
         with self.profiler.profile(hook_name):

pytorch_lightning/trainer/training_loop.py

@@ -595,8 +595,6 @@ def on_train_epoch_end(self, epoch_output: List[List[List[Result]]]) -> None:
             # run training_epoch_end
             # refresh the result for custom logging at the epoch level
             model._current_fx_name = 'training_epoch_end'
-
-            # lightningmodule hook
             training_epoch_end_output = model.training_epoch_end(processed_epoch_output)
 
             if training_epoch_end_output is not None:
@@ -621,7 +619,7 @@ def _on_train_epoch_end_hook(self, processed_epoch_output) -> None:
         hook_name = "on_train_epoch_end"
 
         # set hook_name to model + reset Result obj
-        skip = self.trainer._reset_result_and_set_hook_fx_name(hook_name)
+        skip = self.trainer._reset_result_and_set_fx_name(hook_name)
 
         # always profile hooks
         with self.trainer.profiler.profile(hook_name):

tests/trainer/logging_/test_logger_connector.py

@@ -372,19 +372,17 @@ def test_call_back_validator(tmpdir):
             and func_name not in ["on_train_end", "on_test_end", "on_validation_end"]
         )
         if allowed:
-            validator.check_logging_in_callbacks(current_hook_fx_name=func_name, on_step=on_step, on_epoch=on_epoch)
+            validator.check_logging_in_callbacks(fx_name=func_name, on_step=on_step, on_epoch=on_epoch)
             if not is_start and is_stage:
                 with pytest.raises(MisconfigurationException, match="function supports only"):
-                    validator.check_logging_in_callbacks(
-                        current_hook_fx_name=func_name, on_step=True, on_epoch=on_epoch
-                    )
+                    validator.check_logging_in_callbacks(fx_name=func_name, on_step=True, on_epoch=on_epoch)
         else:
             assert func_name in not_supported
             with pytest.raises(MisconfigurationException, match="function doesn't support"):
-                validator.check_logging_in_callbacks(current_hook_fx_name=func_name, on_step=on_step, on_epoch=on_epoch)
+                validator.check_logging_in_callbacks(fx_name=func_name, on_step=on_step, on_epoch=on_epoch)
 
         # should not fail
-        validator.check_logging_in_callbacks(current_hook_fx_name=None, on_step=None, on_epoch=None)
+        validator.check_logging_in_callbacks(fx_name=None, on_step=None, on_epoch=None)
 
 
 @RunIf(min_gpus=2)

tests/trainer/test_trainer.py

@@ -1912,30 +1912,19 @@ def test_on_load_checkpoint_missing_callbacks(tmpdir):
 
 
 def test_module_current_fx_attributes_reset(tmpdir):
-    """ Ensure that lightning module's attributes related to current hook fx are reset at the end of execution. """
+    """ Ensure that lightning module's attributes related to current fx are reset at the end of execution. """
     model = BoringModel()
-    model.validation_step = None
-    model.training_epoch_end = None
     trainer = Trainer(
         default_root_dir=tmpdir,
-        max_epochs=1,
+        fast_dev_run=1,
         checkpoint_callback=False,
         logger=False,
-        limit_val_batches=0,
     )
+
     trainer.fit(model)
-    assert model._current_fx_name == "", f"_current_fx_name not reset after fit: {model._current_fx_name}"
-    assert (
-        model._current_hook_fx_name is None
-    ), f"_current_hook_fx_name not reset after fit: {model._current_hook_fx_name}"
-    assert (
-        model._current_dataloader_idx is None
-    ), f"_current_dataloader_idx not reset after fit: {model._current_dataloader_idx}"
+    assert model._current_fx_name is None
+    assert model._current_dataloader_idx is None
+
     trainer.test(model)
-    assert model._current_fx_name == "", f"_current_fx_name not reset after test: {model._current_fx_name}"
-    assert (
-        model._current_hook_fx_name is None
-    ), f"_current_hook_fx_name not reset after test: {model._current_hook_fx_name}"
-    assert (
-        model._current_dataloader_idx is None
-    ), f"_current_dataloader_idx not reset after test: {model._current_dataloader_idx}"
+    assert model._current_fx_name is None
+    assert model._current_dataloader_idx is None