def test_device_stats_monitor_logs_for_different_stages(tmp_path):
    """Test that metrics are logged for all stages that is training, testing and validation."""
    model = BoringModel()
    device_stats = DeviceStatsMonitor()

    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=1,
        limit_train_batches=4,
        limit_val_batches=4,
        limit_test_batches=1,
        log_every_n_steps=1,
        accelerator="cpu",
        devices=1,
        callbacks=[device_stats],
        logger=CSVLogger(tmp_path),
        enable_checkpointing=False,
        enable_progress_bar=False,
    )

    # training and validation stages will run
    trainer.fit(model)

    with open(f"{tmp_path}/lightning_logs/version_0/metrics.csv") as csvfile:
        content = csv.reader(csvfile, delimiter=",")
        it = iter(content).__next__()

    # searching for training stage logs
    train_stage_results = [re.match(r".+on_train_batch", i) for i in it]
    train = any(train_stage_results)
    assert train, "training stage logs not found"

    # searching for validation stage logs
    validation_stage_results = [re.match(r".+on_validation_batch", i) for i in it]
    valid = any(validation_stage_results)
    assert valid, "validation stage logs not found"

    # testing stage will run
    trainer.test(model)

    with open(f"{tmp_path}/lightning_logs/version_0/metrics.csv") as csvfile:
        content = csv.reader(csvfile, delimiter=",")
        it = iter(content).__next__()

    # searching for testing stage logs
    test_stage_results = [re.match(r".+on_test_batch", i) for i in it]
    test = any(test_stage_results)

    assert test, "testing stage logs not found"
Test that metrics are logged for all stages that is training, testing and validation.
tests/tests_pytorch/callbacks/test_device_stats_monitor.py