def test_datamodule_hooks_are_profiled(tmp_path):
    """Test that `LightningDataModule` hooks are profiled."""

    def get_trainer():
        return Trainer(
            default_root_dir=tmp_path,
            max_steps=1,
            limit_val_batches=0,
            profiler="simple",
            enable_model_summary=False,
            enable_progress_bar=False,
            logger=False,
        )

    class CustomBoringDataModule(BoringDataModule):
        def state_dict(self):
            return {"temp": 1}

        # override so that it gets called
        def prepare_data(self):
            pass

    model = BoringModel()
    dm = CustomBoringDataModule()
    trainer = get_trainer()
    trainer.fit(model, datamodule=dm)

    profiler = trainer.profiler
    assert isinstance(profiler, SimpleProfiler)

    keys = [
        "[LightningDataModule]CustomBoringDataModule.prepare_data",
        "[LightningDataModule]CustomBoringDataModule.setup",
        "[LightningDataModule]CustomBoringDataModule.state_dict",
        "[LightningDataModule]CustomBoringDataModule.teardown",
    ]
    for key in keys:
        assert key in profiler.recorded_durations
        durations = profiler.recorded_durations[key]
        assert len(durations) == 1
        assert durations[0] > 0

    ckpt_path = trainer.checkpoint_callback.best_model_path
    trainer = get_trainer()
    trainer.fit(model, datamodule=dm, ckpt_path=ckpt_path)
    profiler = trainer.profiler

    keys = [
        "[LightningDataModule]CustomBoringDataModule.prepare_data",
        "[LightningDataModule]CustomBoringDataModule.setup",
        "[LightningDataModule]CustomBoringDataModule.load_state_dict",
        "[LightningDataModule]CustomBoringDataModule.teardown",
    ]
    for key in keys:
        assert key in profiler.recorded_durations
        durations = profiler.recorded_durations[key]
        assert len(durations) == 1
        assert durations[0] > 0
Test that `LightningDataModule` hooks are profiled.
tests/tests_pytorch/core/test_datamodules.py