text: # Copyright The PyTorch Lightning team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import inspect
  import os
  from unittest import mock
  from unittest.mock import PropertyMock
  
  import pytest
  import torch
  
  from pytorch_lightning import Callback, Trainer
  from pytorch_lightning.trainer.states import TrainerState
  from tests.helpers import BoringModel, RandomDataset
  
  
  @pytest.mark.parametrize('max_steps', [1, 2, 3])
  def test_on_before_zero_grad_called(tmpdir, max_steps):
  
      class CurrentTestModel(BoringModel):
          on_before_zero_grad_called = 0
  
          def on_before_zero_grad(self, optimizer):
              self.on_before_zero_grad_called += 1
  
      model = CurrentTestModel()
  
      trainer = Trainer(
          default_root_dir=tmpdir,
          max_steps=max_steps,
          max_epochs=2,
      )
      assert 0 == model.on_before_zero_grad_called
      trainer.fit(model)
      assert max_steps == model.on_before_zero_grad_called
  
      model.on_before_zero_grad_called = 0
      trainer.test(model)
      assert 0 == model.on_before_zero_grad_called
  
  
  def test_training_epoch_end_metrics_collection(tmpdir):
      """ Test that progress bar metrics also get collected at the end of an epoch. """
      num_epochs = 3
  
      class CurrentModel(BoringModel):
  
          def training_step(self, *args, **kwargs):
              output = super().training_step(*args, **kwargs)
              self.log_dict({'step_metric': torch.tensor(-1), 'shared_metric': 100}, logger=False, prog_bar=True)
              return output
  
          def training_epoch_end(self, outputs):
              epoch = self.current_epoch
              # both scalar tensors and Python numbers are accepted
              self.log_dict(
                  {
                      f'epoch_metric_{epoch}': torch.tensor(epoch),
                      'shared_metric': 111
                  },
                  logger=False,
                  prog_bar=True,
              )
  
      model = CurrentModel()
      trainer = Trainer(
          max_epochs=num_epochs,
          default_root_dir=tmpdir,
          overfit_batches=2,
      )
      trainer.fit(model)
      assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
      metrics = trainer.progress_bar_dict
  
      # metrics added in training step should be unchanged by epoch end method
      assert metrics['step_metric'] == -1
      # a metric shared in both methods gets overwritten by epoch_end
      assert metrics['shared_metric'] == 111
      # metrics are kept after each epoch
      for i in range(num_epochs):
          assert metrics[f'epoch_metric_{i}'] == i
  
  
  def test_training_epoch_end_metrics_collection_on_override(tmpdir):
      """ Test that batch end metrics are collected when training_epoch_end is overridden at the end of an epoch. """
  
      class LoggingCallback(Callback):
  
          def on_train_epoch_start(self, trainer, pl_module):
              self.len_outputs = 0
  
          def on_train_epoch_end(self, trainer, pl_module, outputs):
              self.len_outputs = len(outputs[0])
  
      class OverriddenModel(BoringModel):
  
          def on_train_epoch_start(self):
              self.num_train_batches = 0
  
          def training_epoch_end(self, outputs):  # Overridden
              return
  
          def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
              self.num_train_batches += 1
  
      class NotOverriddenModel(BoringModel):
  
          def on_train_epoch_start(self):
              self.num_train_batches = 0
  
          def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
              self.num_train_batches += 1
  
      overridden_model = OverriddenModel()
      not_overridden_model = NotOverriddenModel()
      not_overridden_model.training_epoch_end = None
  
      callback = LoggingCallback()
      trainer = Trainer(
          max_epochs=1,
          default_root_dir=tmpdir,
          overfit_batches=2,
          callbacks=[callback],
      )
  
      trainer.fit(overridden_model)
      # outputs from on_train_batch_end should be accessible in on_train_epoch_end hook
      # if training_epoch_end is overridden
      assert callback.len_outputs == overridden_model.num_train_batches
  
      trainer.fit(not_overridden_model)
      # outputs from on_train_batch_end should be empty
      assert callback.len_outputs == 0
  
  
  @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
  @mock.patch("pytorch_lightning.accelerators.accelerator.Accelerator.lightning_module", new_callable=PropertyMock)
  def test_transfer_batch_hook(model_getter_mock):
  
      class CustomBatch:
  
          def __init__(self, data):
              self.samples = data[0]
              self.targets = data[1]
  
      class CurrentTestModel(BoringModel):
  
          hook_called = False
  
          def transfer_batch_to_device(self, data, device):
              self.hook_called = True
              if isinstance(data, CustomBatch):
                  data.samples = data.samples.to(device)
                  data.targets = data.targets.to(device)
              else:
                  data = super().transfer_batch_to_device(data, device)
              return data
  
      model = CurrentTestModel()
      batch = CustomBatch((torch.zeros(5, 32), torch.ones(5, 1, dtype=torch.long)))
  
      trainer = Trainer(gpus=1)
      # running .fit() would require us to implement custom data loaders, we mock the model reference instead
      model_getter_mock.return_value = model
      batch_gpu = trainer.accelerator_backend.batch_to_device(batch, torch.device('cuda:0'))
      expected = torch.device('cuda', 0)
      assert model.hook_called
      assert batch_gpu.samples.device == batch_gpu.targets.device == expected
  
  
  @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
  @pytest.mark.skipif(
      not os.getenv("PL_RUNNING_SPECIAL_TESTS", '0') == '1', reason="test should be run outside of pytest"
  )
  def test_transfer_batch_hook_ddp(tmpdir):
      """
      Test custom data are properly moved to the right device using ddp
      """
  
      class CustomBatch:
  
          def __init__(self, data):
              self.samples = data[0]
  
          def to(self, device, **kwargs):
              self.samples = self.samples.to(device, **kwargs)
              return self
  
      def collate_fn(batch):
          return CustomBatch(batch)
  
      class TestModel(BoringModel):
  
          def training_step(self, batch, batch_idx):
              assert batch.samples.device == self.device
              assert isinstance(batch_idx, int)
  
          def train_dataloader(self):
              return torch.utils.data.DataLoader(RandomDataset(32, 64), collate_fn=collate_fn)
  
      model = TestModel()
      model.validation_step = None
      model.training_epoch_end = None
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_train_batches=2,
          limit_val_batches=0,
          max_epochs=1,
          weights_summary=None,
          accelerator="ddp",
          gpus=2,
      )
      trainer.fit(model)
  
  
  @pytest.mark.parametrize('max_epochs,batch_idx_', [(2, 5), (3, 8), (4, 12)])
  def test_on_train_batch_start_hook(max_epochs, batch_idx_):
  
      class CurrentModel(BoringModel):
  
          def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
              if batch_idx == batch_idx_:
                  return -1
  
      model = CurrentModel()
      trainer = Trainer(max_epochs=max_epochs)
      trainer.fit(model)
      if batch_idx_ > len(model.val_dataloader()) - 1:
          assert trainer.batch_idx == len(model.val_dataloader()) - 1
          assert trainer.global_step == len(model.val_dataloader()) * max_epochs
      else:
          assert trainer.batch_idx == batch_idx_
          assert trainer.global_step == (batch_idx_ + 1) * max_epochs
  
  
  def test_trainer_model_hook_system(tmpdir):
      """Test the hooks system."""
  
      class HookedModel(BoringModel):
  
          def __init__(self):
              super().__init__()
              self.called = []
  
          def on_after_backward(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_after_backward()
  
          def on_before_zero_grad(self, optimizer):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_before_zero_grad(optimizer)
  
          def on_epoch_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_epoch_start()
  
          def on_epoch_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_epoch_end()
  
          def on_fit_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_fit_start()
  
          def on_fit_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_fit_end()
  
          def on_hpc_load(self, checkpoint):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_hpc_load(checkpoint)
  
          def on_hpc_save(self, checkpoint):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_hpc_save(checkpoint)
  
          def on_load_checkpoint(self, checkpoint):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_load_checkpoint(checkpoint)
  
          def on_save_checkpoint(self, checkpoint):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_save_checkpoint(checkpoint)
  
          def on_pretrain_routine_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_pretrain_routine_start()
  
          def on_pretrain_routine_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_pretrain_routine_end()
  
          def on_train_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_train_start()
  
          def on_train_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_train_end()
  
          def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_train_batch_start(batch, batch_idx, dataloader_idx)
  
          def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_train_batch_end(outputs, batch, batch_idx, dataloader_idx)
  
          def on_train_epoch_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_train_epoch_start()
  
          def on_train_epoch_end(self, outputs):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_train_epoch_end(outputs)
  
          def on_validation_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_start()
  
          def on_validation_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_end()
  
          def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_batch_start(batch, batch_idx, dataloader_idx)
  
          def on_validation_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx)
  
          def on_validation_epoch_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_epoch_start()
  
          def on_validation_epoch_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_epoch_end()
  
          def on_test_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_start()
  
          def on_test_batch_start(self, batch, batch_idx, dataloader_idx):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_batch_start(batch, batch_idx, dataloader_idx)
  
          def on_test_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_batch_end(outputs, batch, batch_idx, dataloader_idx)
  
          def on_test_epoch_start(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_epoch_start()
  
          def on_test_epoch_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_epoch_end()
  
          def on_validation_model_eval(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_model_eval()
  
          def on_validation_model_train(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_validation_model_train()
  
          def on_test_model_eval(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_model_eval()
  
          def on_test_model_train(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_model_train()
  
          def on_test_end(self):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().on_test_end()
  
          def teardown(self, stage: str):
              self.called.append(inspect.currentframe().f_code.co_name)
              super().teardown(stage)
  
      model = HookedModel()
  
      assert model.called == []
  
      # fit model
      trainer = Trainer(
          default_root_dir=tmpdir,
          max_epochs=1,
          limit_val_batches=1,
          limit_train_batches=2,
          limit_test_batches=1,
          progress_bar_refresh_rate=0,
      )
  
      assert model.called == []
  
      trainer.fit(model)
  
      expected = [
          'on_fit_start',
          'on_pretrain_routine_start',
          'on_pretrain_routine_end',
          'on_validation_model_eval',
          'on_validation_start',
          'on_validation_epoch_start',
          'on_validation_batch_start',
          'on_validation_batch_end',
          'on_validation_epoch_end',
          'on_epoch_end',
          'on_validation_end',
          'on_validation_model_train',
          'on_train_start',
          'on_epoch_start',
          'on_train_epoch_start',
          'on_train_batch_start',
          'on_after_backward',
          'on_before_zero_grad',
          'on_train_batch_end',
          'on_train_batch_start',
          'on_after_backward',
          'on_before_zero_grad',
          'on_train_batch_end',
          'on_train_epoch_end',
          'on_epoch_end',
          'on_validation_model_eval',
          'on_validation_start',
          'on_validation_epoch_start',
          'on_validation_batch_start',
          'on_validation_batch_end',
          'on_validation_epoch_end',
          'on_epoch_end',
          'on_save_checkpoint',
          'on_validation_end',
          'on_validation_model_train',
          'on_train_end',
          'on_fit_end',
          'teardown',
      ]
  
      assert model.called == expected
  
      model2 = HookedModel()
      trainer.test(model2)
  
      expected = [
          'on_fit_start',
          'on_test_model_eval',
          'on_test_start',
          'on_test_epoch_start',
          'on_test_batch_start',
          'on_test_batch_end',
          'on_test_epoch_end',
          'on_epoch_end',
          'on_test_end',
          'on_test_model_train',
          'on_fit_end',
          'teardown',  # for 'fit'
          'teardown',  # for 'test'
      ]
  
      assert model2.called == expected
  