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 os
  from typing import Optional
  
  import pytest
  import torch
  import torch.nn.functional as F
  
  from pytorch_lightning import Callback, seed_everything, Trainer
  from pytorch_lightning.accelerators import CPUAccelerator, IPUAccelerator
  from pytorch_lightning.core.lightning import LightningModule
  from pytorch_lightning.plugins import IPUPlugin, IPUPrecisionPlugin
  from pytorch_lightning.trainer.states import RunningStage
  from pytorch_lightning.trainer.supporters import CombinedLoader
  from pytorch_lightning.utilities import _IPU_AVAILABLE, DeviceType
  from pytorch_lightning.utilities.exceptions import MisconfigurationException
  from tests.helpers.boring_model import BoringModel
  from tests.helpers.datamodules import ClassifDataModule
  from tests.helpers.runif import RunIf
  from tests.helpers.simple_models import ClassificationModel
  
  if _IPU_AVAILABLE:
      import poptorch
  
  
  class IPUModel(BoringModel):
      def training_step(self, batch, batch_idx):
          output = self(batch)
          loss = self.loss(batch, output)
          return loss
  
      def validation_step(self, batch, batch_idx):
          output = self(batch)
          loss = self.loss(batch, output)
          return loss
  
      def test_step(self, batch, batch_idx):
          output = self(batch)
          loss = self.loss(batch, output)
          return loss
  
      def training_epoch_end(self, outputs) -> None:
          pass
  
      def validation_epoch_end(self, outputs) -> None:
          pass
  
      def test_epoch_end(self, outputs) -> None:
          pass
  
  
  class IPUClassificationModel(ClassificationModel):
      def training_step(self, batch, batch_idx):
          x, y = batch
          logits = self(x)
          loss = F.cross_entropy(logits, y)
          return loss
  
      def validation_step(self, batch, batch_idx):
          x, y = batch
          logits = self(x)
          acc = self.accuracy(logits, y)
          return acc
  
      def test_step(self, batch, batch_idx):
          x, y = batch
          logits = self(x)
          acc = self.accuracy(logits, y)
          return acc
  
      def accuracy(self, logits, y):
          # todo (sean): currently IPU poptorch doesn't implicit convert bools to tensor
          # hence we use an explicit calculation for accuracy here. Once fixed in poptorch
          # we can use the accuracy metric.
          acc = torch.sum(torch.eq(torch.argmax(logits, -1), y).to(torch.float32)) / len(y)
          return acc
  
      def validation_epoch_end(self, outputs) -> None:
          self.log("val_acc", torch.stack(outputs).mean())
  
      def test_epoch_end(self, outputs) -> None:
          self.log("test_acc", torch.stack(outputs).mean())
  
  
  @pytest.mark.skipif(_IPU_AVAILABLE, reason="test requires non-IPU machine")
  def test_fail_if_no_ipus(tmpdir):
      with pytest.raises(MisconfigurationException, match="IPU Accelerator requires IPU devices to run"):
          Trainer(default_root_dir=tmpdir, ipus=1)
  
      with pytest.raises(MisconfigurationException, match="IPU Accelerator requires IPU devices to run"):
          Trainer(default_root_dir=tmpdir, ipus=1, accelerator="ipu")
  
  
  @RunIf(ipu=True)
  def test_accelerator_selected(tmpdir):
      trainer = Trainer(default_root_dir=tmpdir, ipus=1)
      assert isinstance(trainer.accelerator, IPUAccelerator)
      trainer = Trainer(default_root_dir=tmpdir, ipus=1, accelerator="ipu")
      assert isinstance(trainer.accelerator, IPUAccelerator)
  
  
  @RunIf(ipu=True)
  def test_warning_if_ipus_not_used(tmpdir):
      with pytest.warns(UserWarning, match="IPU available but not used. Set the `ipus` flag in your trainer"):
          Trainer(default_root_dir=tmpdir)
  
  
  @RunIf(ipu=True)
  def test_no_warning_plugin(tmpdir):
      with pytest.warns(None) as record:
          Trainer(default_root_dir=tmpdir, strategy=IPUPlugin(training_opts=poptorch.Options()))
      assert len(record) == 0
  
  
  @RunIf(ipu=True)
  @pytest.mark.parametrize("ipus", [1, 4])
  def test_all_stages(tmpdir, ipus):
      model = IPUModel()
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, ipus=ipus)
      trainer.fit(model)
      trainer.validate(model)
      trainer.test(model)
      trainer.predict(model)
  
  
  @RunIf(ipu=True)
  @pytest.mark.parametrize("ipus", [1, 4])
  def test_inference_only(tmpdir, ipus):
      model = IPUModel()
  
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, ipus=ipus)
      trainer.validate(model)
      trainer.test(model)
      trainer.predict(model)
  
  
  @RunIf(ipu=True)
  def test_optimization(tmpdir):
      seed_everything(42)
  
      dm = ClassifDataModule(length=1024)
      model = IPUClassificationModel()
  
      trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, ipus=2)
  
      # fit model
      trainer.fit(model, dm)
      assert trainer.state.finished, f"Training failed with {trainer.state}"
      assert dm.trainer is not None
  
      # validate
      result = trainer.validate(datamodule=dm)
      assert dm.trainer is not None
      assert result[0]["val_acc"] > 0.7
  
      # test
      result = trainer.test(model, datamodule=dm)
      assert dm.trainer is not None
      test_result = result[0]["test_acc"]
      assert test_result > 0.6
  
      # test saved model
      model_path = os.path.join(tmpdir, "model.pt")
      trainer.save_checkpoint(model_path)
  
      model = IPUClassificationModel.load_from_checkpoint(model_path)
  
      trainer = Trainer(default_root_dir=tmpdir, ipus=2)
  
      result = trainer.test(model, datamodule=dm)
      saved_result = result[0]["test_acc"]
      assert saved_result == test_result
  
  
  @RunIf(ipu=True)
  def test_mixed_precision(tmpdir):
      class TestCallback(Callback):
          def setup(self, trainer: Trainer, pl_module: LightningModule, stage: Optional[str] = None) -> None:
              assert trainer.accelerator.model.precision == 16
              raise SystemExit
  
      model = IPUModel()
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, ipus=1, precision=16, callbacks=TestCallback())
      assert isinstance(trainer.accelerator.precision_plugin, IPUPrecisionPlugin)
      assert trainer.accelerator.precision_plugin.precision == 16
      with pytest.raises(SystemExit):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_pure_half_precision(tmpdir):
      class TestCallback(Callback):
          def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
              assert trainer.accelerator.model.precision == 16
              for param in trainer.accelerator.model.parameters():
                  assert param.dtype == torch.float16
              raise SystemExit
  
      model = IPUModel()
      model = model.half()
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, ipus=1, precision=16, callbacks=TestCallback())
  
      assert isinstance(trainer.accelerator.training_type_plugin, IPUPlugin)
      assert isinstance(trainer.accelerator.precision_plugin, IPUPrecisionPlugin)
      assert trainer.accelerator.precision_plugin.precision == 16
  
      with pytest.raises(SystemExit):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_device_iterations_ipu_plugin(tmpdir):
      class TestCallback(Callback):
          def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
              assert trainer.accelerator.training_type_plugin.device_iterations == 2
              # assert device iterations has been set correctly within the poptorch options
              poptorch_model = trainer.accelerator.training_type_plugin.poptorch_models[RunningStage.TRAINING]
              assert poptorch_model._options.toDict()["device_iterations"] == 2
              raise SystemExit
  
      model = IPUModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          fast_dev_run=True,
          ipus=1,
          strategy=IPUPlugin(device_iterations=2),
          callbacks=TestCallback(),
      )
      assert isinstance(trainer.accelerator.training_type_plugin, IPUPlugin)
      with pytest.raises(SystemExit):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_accumulated_batches(tmpdir):
      class TestCallback(Callback):
          def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
              # ensure the accumulation_scheduler is overridden to accumulate every batch
              # since ipu handle accumulation
              assert trainer.accumulation_scheduler.scheduling == {0: 1}
              # assert poptorch option have been set correctly
              poptorch_model = trainer.accelerator.training_type_plugin.poptorch_models[RunningStage.TRAINING]
              assert poptorch_model._options.Training.toDict()["gradient_accumulation"] == 2
              raise SystemExit
  
      model = IPUModel()
      trainer = Trainer(
          default_root_dir=tmpdir, fast_dev_run=True, ipus=1, accumulate_grad_batches=2, callbacks=TestCallback()
      )
      with pytest.raises(SystemExit):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_stages_correct(tmpdir):
      """Ensure all stages correctly are traced correctly by asserting the output for each stage."""
  
      class StageModel(IPUModel):
          def training_step(self, batch, batch_idx):
              loss = super().training_step(batch, batch_idx)
              # tracing requires a loss value that depends on the model.
              # force it to be a value but ensure we use the loss.
              return (loss - loss) + torch.tensor(1)
  
          def validation_step(self, batch, batch_idx):
              loss = super().validation_step(batch, batch_idx)
              return (loss - loss) + torch.tensor(2)
  
          def test_step(self, batch, batch_idx):
              loss = super().validation_step(batch, batch_idx)
              return (loss - loss) + torch.tensor(3)
  
          def predict_step(self, batch, batch_idx, dataloader_idx=None):
              output = super().predict_step(batch, batch_idx)
              return (output - output) + torch.tensor(4)
  
      class TestCallback(Callback):
          def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) -> None:
              assert outputs["loss"].item() == 1
  
          def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) -> None:
              assert outputs.item() == 2
  
          def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) -> None:
              assert outputs.item() == 3
  
          def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) -> None:
              assert torch.all(outputs == 4).item()
  
      model = StageModel()
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, ipus=1, callbacks=TestCallback())
      trainer.fit(model)
      trainer.test(model)
      trainer.validate(model)
      trainer.predict(model, model.test_dataloader())
  
  
  @RunIf(ipu=True)
  def test_different_accumulate_grad_batches_fails(tmpdir):
      model = IPUModel()
      trainer = Trainer(default_root_dir=tmpdir, ipus=1, accumulate_grad_batches={1: 2})
      with pytest.raises(
          MisconfigurationException, match="IPUs currently does not support different `accumulate_grad_batches`"
      ):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_clip_gradients_fails(tmpdir):
      model = IPUModel()
      trainer = Trainer(default_root_dir=tmpdir, ipus=1, gradient_clip_val=10)
      with pytest.raises(MisconfigurationException, match="IPUs currently do not support clipping gradients."):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_autoreport(tmpdir):
      """Ensure autoreport dumps to a file."""
      model = IPUModel()
      autoreport_path = os.path.join(tmpdir, "report/")
      trainer = Trainer(
          default_root_dir=tmpdir,
          ipus=1,
          fast_dev_run=True,
          strategy=IPUPlugin(autoreport=True, autoreport_dir=autoreport_path),
      )
      trainer.fit(model)
      assert os.path.exists(autoreport_path)
      assert os.path.isfile(autoreport_path + "profile.pop")
  
  
  @RunIf(ipu=True)
  def test_manual_poptorch_opts(tmpdir):
      """Ensure if the user passes manual poptorch Options, we run with the correct object."""
      model = IPUModel()
      inference_opts = poptorch.Options()
      training_opts = poptorch.Options()
  
      trainer = Trainer(
          default_root_dir=tmpdir,
          ipus=1,
          fast_dev_run=True,
          strategy=IPUPlugin(inference_opts=inference_opts, training_opts=training_opts),
      )
      trainer.fit(model)
  
      assert isinstance(trainer.accelerator.training_type_plugin, IPUPlugin)
      assert trainer.accelerator.training_type_plugin.training_opts == training_opts
      assert trainer.accelerator.training_type_plugin.inference_opts == inference_opts
  
  
  @RunIf(ipu=True)
  def test_manual_poptorch_opts_custom(tmpdir):
      """Ensure if the user passes manual poptorch Options with custom parameters set, we respect them in our
      poptorch options and the dataloaders."""
  
      model = IPUModel()
      training_opts = poptorch.Options()
      training_opts.deviceIterations(8)
      training_opts.replicationFactor(2)
      training_opts.Training.gradientAccumulation(2)
  
      inference_opts = poptorch.Options()
      inference_opts.deviceIterations(16)
      inference_opts.replicationFactor(1)
      inference_opts.Training.gradientAccumulation(1)
  
      class TestCallback(Callback):
          def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
              # ensure dataloaders were correctly set up during training.
              plugin = trainer.accelerator.training_type_plugin
              assert isinstance(plugin, IPUPlugin)
              assert plugin.training_opts.replication_factor == 2
              assert plugin.inference_opts.replication_factor == 1
  
              val_dataloader = trainer.val_dataloaders[0]
              train_dataloader = trainer.train_dataloader
              assert isinstance(train_dataloader, CombinedLoader)
              train_dataloader = train_dataloader.loaders
              assert isinstance(val_dataloader, poptorch.DataLoader)
              assert isinstance(train_dataloader, poptorch.DataLoader)
              assert train_dataloader.options.replication_factor == 2
              assert val_dataloader.options.replication_factor == 1
  
      plugin = IPUPlugin(inference_opts=inference_opts, training_opts=training_opts)
      # ensure we default to the training options replication factor
      assert plugin.replication_factor == 2
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy=plugin, callbacks=TestCallback())
      trainer.fit(model)
  
      plugin = trainer.accelerator.training_type_plugin
      assert isinstance(plugin, IPUPlugin)
  
      training_opts = plugin.training_opts
      assert training_opts.device_iterations == 8
      assert training_opts.replication_factor == 2
      assert training_opts.Training.gradient_accumulation == 2
  
      inference_opts = plugin.inference_opts
      assert inference_opts.device_iterations == 16
      assert inference_opts.replication_factor == 1
      assert inference_opts.Training.gradient_accumulation == 1
  
  
  @RunIf(ipu=True)
  def test_replication_factor(tmpdir):
      """Ensure if the user passes manual poptorch Options with custom parameters set, we set them correctly in the
      dataloaders."""
  
      plugin = IPUPlugin()
      trainer = Trainer(ipus=2, default_root_dir=tmpdir, fast_dev_run=True, strategy=plugin)
      assert trainer.ipus == 2
  
  
  @RunIf(ipu=True)
  def test_default_opts(tmpdir):
      """Ensure default opts are set correctly in the IPUPlugin."""
  
      model = IPUModel()
  
      trainer = Trainer(default_root_dir=tmpdir, ipus=1, fast_dev_run=True)
      trainer.fit(model)
      assert isinstance(trainer.accelerator.training_type_plugin, IPUPlugin)
      inference_opts = trainer.accelerator.training_type_plugin.inference_opts
      training_opts = trainer.accelerator.training_type_plugin.training_opts
      for opts in (inference_opts, training_opts):
          assert isinstance(opts, poptorch.Options)
          assert opts.Training.gradient_accumulation == 1
          assert opts.device_iterations == 1
          assert opts.replication_factor == 1
  
  
  @RunIf(ipu=True)
  def test_multi_optimizers_fails(tmpdir):
      """Ensure if there are multiple optimizers, we throw an exception."""
  
      class TestModel(IPUModel):
          def configure_optimizers(self):
              return [torch.optim.Adam(self.parameters()), torch.optim.Adam(self.parameters())]
  
      model = TestModel()
  
      trainer = Trainer(default_root_dir=tmpdir, ipus=1)
      with pytest.raises(MisconfigurationException, match="IPUs currently only support one optimizer."):
          trainer.fit(model)
  
  
  @RunIf(ipu=True)
  def test_precision_plugin(tmpdir):
      """Ensure precision plugin value is set correctly."""
  
      plugin = IPUPrecisionPlugin(precision=16)
      assert plugin.precision == 16
  
  
  @RunIf(ipu=True)
  def test_accelerator_ipu():
  
      trainer = Trainer(accelerator="ipu", ipus=1)
  
      assert trainer._device_type == "ipu"
      assert isinstance(trainer.accelerator, IPUAccelerator)
  
      with pytest.raises(
          MisconfigurationException, match="You passed `accelerator='ipu'`, but you didn't pass `ipus` to `Trainer`"
      ):
          trainer = Trainer(accelerator="ipu")
  
      trainer = Trainer(accelerator="auto", ipus=8)
  
      assert trainer._device_type == "ipu"
      assert isinstance(trainer.accelerator, IPUAccelerator)
  
  
  @RunIf(ipu=True)
  def test_accelerator_cpu_with_ipus_flag():
  
      trainer = Trainer(accelerator="cpu", ipus=1)
  
      assert trainer._device_type == "cpu"
      assert isinstance(trainer.accelerator, CPUAccelerator)
  
  
  @RunIf(ipu=True)
  def test_accelerator_ipu_with_devices():
  
      trainer = Trainer(accelerator="ipu", devices=8)
  
      assert trainer.ipus == 8
      assert isinstance(trainer.training_type_plugin, IPUPlugin)
      assert isinstance(trainer.accelerator, IPUAccelerator)
  
  
  @RunIf(ipu=True)
  def test_accelerator_auto_with_devices_ipu():
  
      trainer = Trainer(accelerator="auto", devices=8)
  
      assert trainer._device_type == "ipu"
      assert trainer.ipus == 8
  
  
  @RunIf(ipu=True)
  def test_accelerator_ipu_with_ipus_priority():
      """Test for checking `ipus` flag takes priority over `devices`."""
  
      ipus = 8
      with pytest.warns(UserWarning, match="The flag `devices=1` will be ignored,"):
          trainer = Trainer(accelerator="ipu", devices=1, ipus=ipus)
  
      assert trainer.ipus == ipus
  
  
  @RunIf(ipu=True)
  def test_set_devices_if_none_ipu():
  
      trainer = Trainer(accelerator="ipu", ipus=8)
      assert trainer.devices == 8
  
  
  @RunIf(ipu=True)
  def test_strategy_choice_ipu_plugin(tmpdir):
      trainer = Trainer(strategy=IPUPlugin(), accelerator="ipu", devices=8)
      assert isinstance(trainer.training_type_plugin, IPUPlugin)
  
  
  @RunIf(ipu=True)
  def test_device_type_when_training_plugin_ipu_passed(tmpdir):
  
      trainer = Trainer(strategy=IPUPlugin(), ipus=8)
      assert isinstance(trainer.training_type_plugin, IPUPlugin)
      assert trainer._device_type == DeviceType.IPU
      assert isinstance(trainer.accelerator, IPUAccelerator)
  