text: #!/usr/bin/env python3
  
  import inspect
  import logging
  
  import pytorch_lightning as pl
  import torch
  from reagent.core.tensorboardX import SummaryWriterContext
  from reagent.core.utils import lazy_property
  from typing_extensions import final
  
  
  logger = logging.getLogger(__name__)
  
  
  class ReAgentLightningModule(pl.LightningModule):
      def __init__(self, automatic_optimization=True):
          super().__init__()
          self._automatic_optimization = automatic_optimization
          self._training_step_generator = None
          self._reporter = pl.loggers.base.DummyExperiment()
          # For the generator API
          self._verified_steps = False
          # For summary_writer property
          self._summary_writer_logger = None
          self._summary_writer = None
          # To enable incremental training
          self.register_buffer("_next_stopping_epoch", None)
          self.register_buffer("_cleanly_stopped", None)
          self._next_stopping_epoch = torch.tensor([-1]).int()
          self._cleanly_stopped = torch.ones(1)
          self._setup_input_type()
          self.batches_processed_this_epoch = 0
          self.all_batches_processed = 0
  
      def _setup_input_type(self):
          self._training_batch_type = None
          sig = inspect.signature(self.train_step_gen)
          assert "training_batch" in sig.parameters
          param = sig.parameters["training_batch"]
          annotation = param.annotation
          if annotation == inspect.Parameter.empty:
              return
          if hasattr(annotation, "from_dict"):
              self._training_batch_type = annotation
  
      def set_reporter(self, reporter):
          if reporter is None:
              reporter = pl.loggers.base.DummyExperiment()
          self._reporter = reporter
          return self
  
      @property
      def reporter(self):
          return self._reporter
  
      def set_clean_stop(self, clean_stop: bool):
          if clean_stop:
              self._cleanly_stopped = torch.ones(1)
          else:
              self._cleanly_stopped = torch.zeros(1)
  
      def increase_next_stopping_epochs(self, num_epochs: int):
          self._next_stopping_epoch += num_epochs
          self.set_clean_stop(False)
          return self
  
      def train_step_gen(self, training_batch, batch_idx: int):
          """
          Implement training step as generator here
          """
          raise NotImplementedError
  
      def soft_update_result(self) -> torch.Tensor:
          """
          A dummy loss to trigger soft-update
          """
          one = torch.ones(1, requires_grad=True)
          return one + one
  
      @property
      def summary_writer(self):
          """
          Accessor to TensorBoard's SummaryWriter
          """
          if self._summary_writer_logger is self.logger:
              # If self.logger doesn't change between call, then return cached result
              return self._summary_writer
  
          # Invalidate
          self._summary_writer = None
          self._summary_writer_logger = self.logger
  
          if isinstance(self.logger, pl.loggers.base.LoggerCollection):
              for logger in self.logger._logger_iterable:
                  if isinstance(logger, pl.loggers.tensorboard.TensorBoardLogger):
                      self._summary_writer = logger.experiment
                      break
          elif isinstance(logger, pl.loggers.tensorboard.TensorBoardLogger):
              self._summary_writer = logger.experiment
  
          return self._summary_writer
  
      # pyre-fixme[14]: `training_step` overrides method defined in `LightningModule`
      #  inconsistently.
      def training_step(self, batch, batch_idx: int, optimizer_idx: int = 0):
          assert (optimizer_idx == 0) or (self._num_optimizing_steps > 1)
  
          if self._training_step_generator is None:
              if self._training_batch_type and isinstance(batch, dict):
                  batch = self._training_batch_type.from_dict(batch)
              self._training_step_generator = self.train_step_gen(batch, batch_idx)
  
          ret = next(self._training_step_generator)
  
          if optimizer_idx == self._num_optimizing_steps - 1:
              if not self._verified_steps:
                  try:
                      next(self._training_step_generator)
                  except StopIteration:
                      self._verified_steps = True
                  if not self._verified_steps:
                      raise RuntimeError(
                          "training_step_gen() yields too many times."
                          "The number of yields should match the number of optimizers,"
                          f" in this case {self._num_optimizing_steps}"
                      )
              self._training_step_generator = None
              SummaryWriterContext.increase_global_step()
  
          return ret
  
      def optimizers(self, use_pl_optimizer: bool = True):
          o = super().optimizers(use_pl_optimizer)
          if isinstance(o, list):
              return o
          return [o]
  
      @lazy_property
      def _num_optimizing_steps(self) -> int:
          return len(self.configure_optimizers())
  
      @final
      def on_epoch_end(self):
          logger.info(
              f"Finished epoch with {self.batches_processed_this_epoch} batches processed"
          )
          self.batches_processed_this_epoch = 0
          # Flush the reporter which has accumulated data in
          # training/validation/test
          self.reporter.flush(self.current_epoch)
  
          # Tell the trainer to stop.
          if self.current_epoch == self._next_stopping_epoch.item():
              self.trainer.should_stop = True
  
      @final
      def on_train_batch_end(self, *args, **kwargs):
          self.batches_processed_this_epoch += 1
          self.all_batches_processed += 1
  
      @final
      def on_validation_batch_end(self, *args, **kwargs):
          self.batches_processed_this_epoch += 1
  
      @final
      def on_test_batch_end(self, *args, **kwargs):
          self.batches_processed_this_epoch += 1
  
      def train(self, *args):
          # trainer.train(batch) was the old, pre-Lightning ReAgent trainer API.
          # make sure that nobody is trying to call trainer.train() this way.
          # trainer.train() or trainer.train(True/False) is allowed - this puts the network into training/eval mode.
          if (len(args) == 0) or ((len(args) == 1) and (isinstance(args[0], bool))):
              super().train(*args)
          else:
              raise NotImplementedError(
                  "Method .train() is not used for ReAgent Lightning trainers. Please use .fit() method of the pl.Trainer instead"
              )
  
  
  class StoppingEpochCallback(pl.Callback):
      """
      We use this callback to control the number of training epochs in incremental
      training. Epoch & step counts are not reset in the checkpoint. If we were to set
      `max_epochs` on the trainer, we would have to keep track of the previous `max_epochs`
      and add to it manually. This keeps the infomation in one place.
  
      Note that we need to set `_cleanly_stopped` back to True before saving the checkpoint.
      This is done in `ModelManager.save_trainer()`.
      """
  
      def __init__(self, num_epochs):
          super().__init__()
          self.num_epochs = num_epochs
  
      def on_pretrain_routine_end(self, trainer, pl_module):
          assert isinstance(pl_module, ReAgentLightningModule)
          cleanly_stopped = pl_module._cleanly_stopped.item()
          logger.info(f"cleanly stopped: {cleanly_stopped}")
          if cleanly_stopped:
              pl_module.increase_next_stopping_epochs(self.num_epochs)
  
  
  def has_test_step_override(trainer_module: ReAgentLightningModule):
      """Detect if a subclass of LightningModule has test_step overridden"""
      return type(trainer_module).test_step != pl.LightningModule.test_step
  