Fix self.log(on_epoch=True) on_batch_start

[carmocca]

What does this PR do?

Minimal repro:

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer

from pytorch_lightning.callbacks import Callback


class ACallback(Callback):

    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        print(trainer.callback_metrics)

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        pl_module.log("on_train_batch_start", 1.0, reduce_fx="sum")

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        pl_module.log("on_train_batch_end", 1.0, reduce_fx="sum")


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        return {"loss": loss}

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        progress_bar_refresh_rate=0,
        limit_train_batches=3,
        max_epochs=1,
        callbacks=[ACallback()],
        log_every_n_steps=1,
    )
    trainer.fit(model, train_dataloader=train_data)


if __name__ == "__main__":
    run()

which prints

{'on_train_batch_start': tensor(2.), 'on_train_batch_end': tensor(3.)}
# should be
# {'on_train_batch_start': tensor(3.), 'on_train_batch_end': tensor(3.)}

This happens because logger_connector.on_train_split_start is not called before the start hooks run.

Related to #9772

Does your PR introduce any breaking changes? If yes, please list them.

None

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

[codecov]

Codecov Report

Merging #9780 (966f709) into master (db4e770) will decrease coverage by 0%.
The diff coverage is 94%.

@@          Coverage Diff           @@
##           master   #9780   +/-   ##
======================================
- Coverage      93%     93%   -0%     
======================================
  Files         179     179           
  Lines       15833   15827    -6     
======================================
- Hits        14676   14669    -7     
- Misses       1157    1158    +1     

[tchaton]

Hey @carmocca. Mind resolving the conflicts and ping me back for review ?

Best,
T.C

[tchaton]

LGTM !