Running `test` with LightningCLI, the program can quit before the test loop ends

Nobody has claimed this yet.

### Bug description

Within my `LightningModule`, I used `self.log_dict(metrics, on_step=True, on_epoch=True)` in `test_step`, and run with `python main.py test --config config.yaml`, with `main.py` containing only `cli = LightningCLI()`, and `config.yaml` providing both the datasets and model. The `TensorBoardLogger` is used.

However, after the programs ends, sometimes I can normally get the metrics `epoch`, `test_accuracy_epoch` and `test_loss_epoch` in the logger file, but at most attempts these 3 metrics didn't show up, and step-level logged objects can always be seen normally.

When the problems occurs, nothing abnormal can be seen from command line outputs. It looks as if the program quited normally.

I find a walkaround to be sleeping for a while in `main.py` right after `cli = LightningCLI()`. It seems like this is because a child thread is not waited to the end.

### What version are you seeing the problem on?

v2.2

### How to reproduce the bug

main.py

```python
from lightning.pytorch.cli import LightningCLI
from lightning.pytorch.loggers import TensorBoardLogger
from lightning.pytorch.callbacks import ModelCheckpoint

from model import Model
from datamodule import DataModule

def cli_main():
cli = LightningCLI()

if __name__ == "__main__":
cli_main()

from time import sleep
sleep(2)
# The problem can be solved by adding sleep.
```

config.yaml
```yaml
# lightning.pytorch==2.2.5
ckpt_path: null
seed_everything: 0
model:
class_path: model.Model
init_args:
learning_rate: 1e-3
data:
class_path: datamodule.DataModule
init_args:
data_dir: data
trainer:
accelerator: gpu
strategy: auto
devices: 1
num_nodes: 1
precision: null
fast_dev_run: false
max_epochs: 100
min_epochs: null
max_steps: -1
min_steps: null
max_time: null
limit_train_batches: null
limit_val_batches: 10
limit_test_batches: null
limit_predict_batches: null
logger:
class_path: lightning.pytorch.loggers.TensorBoardLogger
init_args:
save_dir: lightning_logs/resnet50
name: normalized
callbacks:
class_path: lightning.pytorch.callbacks.ModelCheckpoint
init_args:
save_top_k: 5
monitor: valid_loss
filename: "{epoch}-{step}-{valid_loss:.8f}"
overfit_batches: 0.0
val_check_interval: 50
check_val_every_n_epoch: 1
num_sanity_val_steps: null
log_every_n_steps: 50
enable_checkpointing: null
enable_progress_bar: null
enable_model_summary: null
accumulate_grad_batches: 1
gradient_clip_val: null
gradient_clip_algorithm: null
deterministic: false
benchmark: null
inference_mode: true
use_distributed_sampler: true
profiler: null
detect_anomaly: false
barebones: false
plugins: null
sync_batchnorm: true
reload_dataloaders_every_n_epochs: 0
default_root_dir: null
```

model.py
```Python
import torch
from torch import nn
import torch.nn.functional as F
import lightning as pl
from torchvision.models import resnet50

class Model(pl.LightningModule):
def __init__(self, learning_rate: float):
super().__init__()

self.save_hyperparameters()

CHARS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
class_num = len(CHARS)
self.text_len = 4

resnet = resnet50()
resnet.conv1 = nn.Conv2d(
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
layers = list(resnet.children())
self.resnet = nn.Sequential(*layers[:9])
self.linear = nn.Linear(512, class_num)
self.softmax = nn.Softmax(2)

def _calc_softmax(self, x: torch.Tensor) -> torch.Tensor:
x = self.resnet(x) # (batch, 2048, 1, 1)
x = x.reshape(x.shape[0], self.text_len, -1) # (batch, 4, 512)
x = self.linear(x) # (batch, 4, 62)
x = self.softmax(x) # (batch, 4, 62)
return x

def forward(self, x: torch.Tensor) -> torch.Tensor:
# in lightning, forward defines the prediction/inference actions
x = self._calc_softmax(x) # (batch, 4, 62)
return torch.argmax(x, 2) # (batch, 4)

def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
# training_step defined the train loop.
# It is independent of forward
img, target = batch
batch_size = img.shape[0]

pred_softmax = self._calc_softmax(img) # (batch, 4, 62)
pred_softmax_permute = pred_softmax.permute((0, 2, 1)) # (batch, 62, 4)

loss = F.cross_entropy(pred_softmax_permute, target)

with torch.no_grad():
pred = torch.argmax(pred_softmax, 2) # (batch, 4)
char_correct = (pred == target).sum(1) # (batch)
batch_correct = (char_correct == self.text_len).sum()
batch_accuracy = batch_correct / batch_size

metrics = {"train_accuracy": batch_accuracy, "train_loss": loss}
self.log_dict(metrics, prog_bar=True, logger=True, on_step=True, on_epoch=True)

return loss

def configure_optimizers(self) -> torch.optim.Optimizer:
optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
return optimizer

def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
# validation_step defined the validation loop.
# It is independent of forward
img, target = batch
batch_size = img.shape[0]

pred_softmax = self._calc_softmax(img) # (batch, 4, 62)
pred_softmax_permute = pred_softmax.permute((0, 2, 1)) # (batch, 62, 4)

loss = F.cross_entropy(pred_softmax_permute, target)

with torch.no_grad():
pred = torch.argmax(pred_softmax, 2) # (batch, 4)
char_correct = (pred == target).sum(1) # (batch)
batch_correct = (char_correct == self.text_len).sum()
batch_accuracy = batch_correct / batch_size

metrics = {"valid_accurary": batch_accuracy, "valid_loss": loss}
self.log_dict(metrics, prog_bar=True, logger=True, on_step=True, on_epoch=True)

return loss

def test_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
# test_step defined the test loop.
# It is independent of forward
img, target = batch
batch_size = img.shape[0]

pred_softmax = self._calc_softmax(img) # (batch, 4, 62)
pred_softmax_permute = pred_softmax.permute((0, 2, 1)) # (batch, 62, 4)

loss = F.cross_entropy(pred_softmax_permute, target)

with torch.no_grad():
pred = torch.argmax(pred_softmax, 2) # (batch, 4)
char_correct = (pred == target).sum(1) # (batch)
batch_correct = (char_correct == self.text_len).sum()
batch_accuracy = batch_correct / batch_size

metrics = {"test_accurary": batch_accuracy, "test_loss": loss}
self.log_dict(metrics, prog_bar=True, logger=True, on_step=True, on_epoch=True)
## The `on_epoch` part of behaviors are unstable, but `test_accuracy_step` can always be seen.
## If `on_step=False` and `on_epoch=True`, it works fine to me.

return loss
```
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
- PyTorch Lightning Version: 2.2.5
- PyTorch Version: 2.3.1+cu121
- Python version: 3.12.4
- OS: Windows 11
- CUDA/cuDNN version: 12.1
- GPU models and configuration: GTX 1650
- How you installed Lightning: pip
```

### More info

_No response_

cc @mauvilsa

Start by running the `python main.py test --config config.yaml` reproduction with the `main.py`, `config.yaml`, and `model.py` examples, focusing on `LightningCLI`, `test_step`, and `TensorBoardLogger`. Confirm whether epoch-level test metrics are consistently written after the test loop exits, without relying on the `sleep(2)` workaround.

Written by the indexing model from the issue text.