Repository: Lightning-AI/pytorch-lightning
Pull request: Add Profiler table kwargs
## What does this PR do?

Adds table_kwargs to the PyTorchProfiler class to expose more formatting options to the user directly.

Fixes #17367

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **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 minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] 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

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


Problem statement:
Expose more arguments for PyTorchProfiler
### Description & Motivation

There are only limited arguments for PyTorch profiler exposed to the user. Some other arguments might be useful, such as `max_name_column_width` when the path is long and truncated. 

### Pitch

Expose more arguments of PyTorch profiler such as `max_name_column_width`. 

### Alternatives

Customize a lightning profiler to get a new PyTorch profiler. 

### Additional context

_No response_

cc @borda @carmocca @nbcsm @guotuofeng

Hints from issue discussion:
The class definition accepts `**profiler_kwargs`. So you should be able to pass this already: https://github.com/Lightning-AI/lightning/blob/a5c43d3b2b80f5fc769d7ed0ea511c0bd6733c6b/src/lightning/pytorch/profilers/pytorch.py#L243
> The class definition accepts **profiler_kwargs. So you should be able to pass this already:

Thanks for taking a look! But I don't think passing arguments like `max_name_column_width` will work. It should be added here similar to `_sort_by_key` and `_row_limit`.

https://github.com/Lightning-AI/lightning/blob/a5c43d3b2b80f5fc769d7ed0ea511c0bd6733c6b/src/lightning/pytorch/profilers/pytorch.py#L488
I see! Then you could add `table_kwargs` to the signature

Commit message: Mark table_kwargs test for Python 3.10 and above

Patch to review:
diff --git a/src/lightning/pytorch/profilers/pytorch.py b/src/lightning/pytorch/profilers/pytorch.py
index 1971646eeb81b..fe3ab1c18968c 100644
--- a/src/lightning/pytorch/profilers/pytorch.py
+++ b/src/lightning/pytorch/profilers/pytorch.py
@@ -21,7 +21,7 @@
 
 import torch
 from torch import nn, Tensor
-from torch.autograd.profiler import record_function
+from torch.autograd.profiler import EventList, record_function
 
 from lightning.fabric.accelerators.cuda import is_cuda_available
 from lightning.pytorch.profilers.profiler import Profiler
@@ -30,7 +30,6 @@
 from lightning.pytorch.utilities.rank_zero import rank_zero_warn, WarningCache
 
 if TYPE_CHECKING:
-    from torch.autograd.profiler import EventList
     from torch.utils.hooks import RemovableHandle
 
     from lightning.pytorch.core.module import LightningModule
@@ -239,6 +238,7 @@ def __init__(
         row_limit: int = 20,
         sort_by_key: Optional[str] = None,
         record_module_names: bool = True,
+        table_kwargs: Optional[Dict[str, Any]] = None,
         **profiler_kwargs: Any,
     ) -> None:
         r"""This profiler uses PyTorch's Autograd Profiler and lets you inspect the cost of.
@@ -279,6 +279,8 @@ def __init__(
 
             record_module_names: Whether to add module names while recording autograd operation.
 
+            table_kwargs: Dictionary with keyword arguments for the summary table.
+
             \**profiler_kwargs: Keyword arguments for the PyTorch profiler. This depends on your PyTorch version
 
         Raises:
@@ -296,6 +298,7 @@ def __init__(
         self._sort_by_key = sort_by_key or f"{'cuda' if profiler_kwargs.get('use_cuda', False) else 'cpu'}_time_total"
         self._record_module_names = record_module_names
         self._profiler_kwargs = profiler_kwargs
+        self._table_kwargs = table_kwargs if table_kwargs is not None else {}
 
         self.profiler: Optional[_PROFILER] = None
         self.function_events: Optional["EventList"] = None
@@ -314,6 +317,19 @@ def __init__(
                 f"Found sort_by_key: {self._sort_by_key}. Should be within {self.AVAILABLE_SORT_KEYS}. "
             )
 
+        for key in self._table_kwargs:
+            if key in {"sort_by", "row_limit"}:
+                raise KeyError(
+                    f"Found invalid table_kwargs key: {key}. This is already a positional argument of the Profiler."
+                )
+            valid_table_keys = set(inspect.signature(EventList.table).parameters.keys()) - {
+                "self",
+                "sort_by",
+                "row_limit",
+            }
+            if key not in valid_table_keys:
+                raise KeyError(f"Found invalid table_kwargs key: {key}. Should be within {valid_table_keys}.")
+
     def _init_kineto(self, profiler_kwargs: Any) -> None:
         has_schedule = "schedule" in profiler_kwargs
         self._has_on_trace_ready = "on_trace_ready" in profiler_kwargs
@@ -485,7 +501,7 @@ def summary(self) -> str:
             self.function_events.export_chrome_trace(path_to_trace)
 
         data = self.function_events.key_averages(group_by_input_shapes=self._group_by_input_shapes)
-        table = data.table(sort_by=self._sort_by_key, row_limit=self._row_limit)
+        table = data.table(sort_by=self._sort_by_key, row_limit=self._row_limit, **self._table_kwargs)
 
         recorded_stats = {"records": table}
         return self._stats_to_str(recorded_stats)
diff --git a/tests/tests_pytorch/profilers/test_profiler.py b/tests/tests_pytorch/profilers/test_profiler.py
index e5c9098e3f760..3fa427cf98331 100644
--- a/tests/tests_pytorch/profilers/test_profiler.py
+++ b/tests/tests_pytorch/profilers/test_profiler.py
@@ -14,6 +14,7 @@
 import logging
 import os
 import platform
+import sys
 import time
 from copy import deepcopy
 from unittest.mock import patch
@@ -619,3 +620,37 @@ def test_profile_callbacks(tmpdir):
         e.name == "[pl][profile][Callback]EarlyStopping{'monitor': 'train_loss', 'mode': 'min'}.on_validation_start"
         for e in pytorch_profiler.function_events
     )
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10 or higher")
+def test_profiler_table_kwargs_summary_length(tmpdir):
+    """Test if setting max_name_column_width in table_kwargs changes summary length."""
+
+    summaries = []
+    # Default table_kwargs (None) sets max_name_column_width to 55
+    for table_kwargs in [None, {"max_name_column_width": 20}, {"max_name_column_width": 100}]:
+        pytorch_profiler = PyTorchProfiler(dirpath=tmpdir, filename="profile", table_kwargs=table_kwargs)
+        model = BoringModel()
+        trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, fast_dev_run=1, profiler=pytorch_profiler)
+        trainer.fit(model)
+        summaries.append(pytorch_profiler.summary())
+
+    # Check if setting max_name_column_width results in a different number of characters in the summary
+    assert len(summaries[0]) > len(summaries[1])
+    assert len(summaries[0]) < len(summaries[2])
+
+
+def test_profiler_invalid_table_kwargs(tmpdir):
+    """Test if passing invalid keyword arguments raise expected error."""
+
+    for key in {"row_limit", "sort_by"}:
+        with pytest.raises(
+            KeyError,
+            match=f"Found invalid table_kwargs key: {key}. This is already a positional argument of the Profiler.",
+        ):
+            PyTorchProfiler(table_kwargs={key: None}, dirpath=tmpdir, filename="profile")
+
+    for key in {"self", "non_existent_keyword_arg"}:
+        with pytest.raises(KeyError) as exc_info:
+            PyTorchProfiler(table_kwargs={key: None}, dirpath=tmpdir, filename="profile")
+        assert exc_info.value.args[0].startswith(f"Found invalid table_kwargs key: {key}.")


Code review comments:
diff_hunk: @@ -619,3 +620,37 @@ def test_profile_callbacks(tmpdir):
         e.name == "[pl][profile][Callback]EarlyStopping{'monitor': 'train_loss', 'mode': 'min'}.on_validation_start"
         for e in pytorch_profiler.function_events
     )
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10 or higher")
path: tests/tests_pytorch/profilers/test_profiler.py
text: @user2:
```suggestion
@user1(min_python="3.10")
```

diff_hunk: @@ -619,3 +620,37 @@ def test_profile_callbacks(tmpdir):
         e.name == "[pl][profile][Callback]EarlyStopping{'monitor': 'train_loss', 'mode': 'min'}.on_validation_start"
         for e in pytorch_profiler.function_events
     )
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10 or higher")
+def test_profiler_table_kwargs_summary_length(tmpdir):
+    """Test if setting max_name_column_width in table_kwargs changes summary length."""
+
+    summaries = []
+    # Default table_kwargs (None) sets max_name_column_width to 55
+    for table_kwargs in [None, {"max_name_column_width": 20}, {"max_name_column_width": 100}]:
+        pytorch_profiler = PyTorchProfiler(dirpath=tmpdir, filename="profile", table_kwargs=table_kwargs)
+        model = BoringModel()
+        trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, fast_dev_run=1, profiler=pytorch_profiler)
+        trainer.fit(model)
path: tests/tests_pytorch/profilers/test_profiler.py
text: @user1:
Can you add data to the profile manually instead of starting up a Trainer just to get some data in?