Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/trainer.py (Python)

Code before:
    ):
        r"""
        Runs the full optimization routine.

        Args:
            datamodule: A instance of :class:`LightningDataModule`.

            model: Model to fit.

            train_dataloader: A Pytorch DataLoader with training samples. If the model has
                a predefined train_dataloader method this will be skipped.

            val_dataloaders: Either a single Pytorch Dataloader or a list of them, specifying validation samples.
                If the model has a predefined val_dataloaders method this will be skipped

        """
        # bookkeeping
        self._state = TrainerState.RUNNING

        # ----------------------------
        # SET UP TRAINING
        # ----------------------------
        self.accelerator_backend = self.accelerator_connector.select_accelerator()
        self.local_rank = self.accelerator_backend.cluster_environment.local_rank()
        self.node_rank = self.accelerator_backend.cluster_environment.node_rank()
        self.global_rank = self.accelerator_backend.cluster_environment.global_rank()
        self.accelerator_backend.setup(model)

        # ----------------------------
        # LINK DATA
        # ----------------------------
        # setup data, etc...
        self.train_loop.setup_fit(model, train_dataloader, val_dataloaders, datamodule)

        # hook
        self.data_connector.prepare_data(model)

        # bookkeeping
        # we reuse fit in .test() but change its behavior using this flag
        self.testing = os.environ.get('PL_TESTING_MODE', self.testing)

        # ----------------------------
        # INSPECT THESE FOR MAIN LOOPS
        # ----------------------------
        # assign training and eval functions... inspect these to see the train and eval loops :)
        self.accelerator_backend.train_loop = self.train
        self.accelerator_backend.validation_loop = self.run_evaluation
        self.accelerator_backend.test_loop = self.run_evaluation

        # ----------------------------
        # TRAIN

Review comment (suggestion):
```suggestion
        self.accelerator_backend.cluster_environment.set_ranks_to_trainer(self)
```

Code after:
        self.logger_connector.on_trainer_init(
            logger,
            flush_logs_every_n_steps,
            log_every_n_steps,
            move_metrics_to_cpu,
        )

        # init debugging flags
        self.debugging_connector.on_init_start(
            limit_train_batches,
            limit_val_batches,
            limit_test_batches,
            limit_predict_batches,
            val_check_interval,
            overfit_batches,
            fast_dev_run,
        )

        # Callback system
        self.on_init_end()

    def _setup_on_init(
        self,
        num_sanity_val_steps: int,
    ) -> None:
        self._log_device_info()

        self.should_stop = False
        self.state = TrainerState()
        self.num_training_batches = 0
        self.train_dataloader = None

        if num_sanity_val_steps == -1:
            self.num_sanity_val_steps = float("inf")
        else:
            self.num_sanity_val_steps = num_sanity_val_steps

        self.num_sanity_val_batches = []
        self.num_test_batches = []
        self.num_val_batches = []
        self.test_dataloaders = None
        self.val_dataloaders = None

        # .validate() and .test() set this when they load a checkpoint
        self.validated_ckpt_path = None
        self.tested_ckpt_path = None

        # when true, print evaluation results in .validate() and .test()
        self.verbose_evaluate = True

        self.num_predict_batches = []