Profile dataloader initialization and update profiler docs

Author: yuxin00j

docs/source-pytorch/tuning/profiler_basic.rst

@@ -80,7 +80,7 @@ Once the **.fit()** function has completed, you'll see an output like this:
 
     Profiler Report
 
-    Profile stats for: get_train_batch
+    Profile stats for: setup_train_dataloader
             4869394 function calls (4863767 primitive calls) in 18.893 seconds
     Ordered by: cumulative time
     List reduced from 76 to 10 due to restriction <10>

src/lightning/pytorch/CHANGELOG.md

@@ -24,6 +24,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed PyTorch Lightning profiler not capturing dataloader worker initialization time ([#21771](https://github.com/Lightning-AI/pytorch-lightning/issues/21771))
+
 - Fixed non-zero process exits in `CombinedLoader.reset()` with large tensors and persistent spawned workers by avoiding explicit `_shutdown_workers()` calls and relying on iterator cleanup via `del` [#21708](https://github.com/Lightning-AI/pytorch-lightning/issues/21708)
 
 - Fixed `SIGTERMException` producing a zero exit code instead of 143 (128 + SIGTERM) ([#21623](https://github.com/Lightning-AI/pytorch-lightning/issues/21623))

src/lightning/pytorch/loops/fit_loop.py

@@ -272,7 +272,8 @@ def setup_data(self) -> None:
 
         self._data_fetcher = _select_data_fetcher(trainer, RunningStage.TRAINING)
         self._data_fetcher.setup(combined_loader)
-        iter(self._data_fetcher)  # creates the iterator inside the fetcher
+        with trainer.profiler.profile("setup_train_dataloader"):
+            iter(self._data_fetcher)  # creates the iterator inside the fetcher
         max_batches = sized_len(combined_loader)
         self.max_batches = max_batches if max_batches is not None else float("inf")
         has_len_all_ranks_ = has_len_all_ranks(combined_loader, trainer.strategy, allow_zero_length)

tests/tests_pytorch/profilers/test_profiler.py

@@ -731,3 +731,14 @@ def test_profiler_invalid_table_kwargs(tmp_path):
         with pytest.raises(KeyError) as exc_info:
             PyTorchProfiler(table_kwargs={key: None}, dirpath=tmp_path, filename="profile")
         assert exc_info.value.args[0].startswith(f"Found invalid table_kwargs key: {key}.")
+
+
+@skip_advanced_profiler_py312
+def test_setup_train_dataloader_profiled_actions(tmp_path):
+    """Ensure that the 'setup_train_dataloader' action is successfully recorded in the profiler."""
+    profiler = AdvancedProfiler(dirpath=tmp_path, filename="profiler")
+    model = BoringModel()
+    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=2, profiler=profiler)
+    trainer.fit(model)
+
+    assert "setup_train_dataloader" in profiler.profiled_actions