Rename profiler to profilers (Lightning-AI#12308)

Co-authored-by: Carlos Mocholí <[email protected]>
Co-authored-by: Akihiro Nitta <[email protected]>

.github/CODEOWNERS

@@ -33,8 +33,8 @@
 /src/pytorch_lightning/loops                @tchaton @awaelchli @justusschock @carmocca
 /src/pytorch_lightning/overrides            @tchaton @SeanNaren @borda
 /src/pytorch_lightning/plugins              @tchaton @SeanNaren @awaelchli @justusschock
-/src/pytorch_lightning/profiler             @williamfalcon @tchaton @borda @carmocca
-/src/pytorch_lightning/profiler/pytorch.py  @nbcsm @guotuofeng
+/src/pytorch_lightning/profilers             @williamfalcon @tchaton @borda @carmocca
+/src/pytorch_lightning/profilers/pytorch.py  @nbcsm @guotuofeng
 /src/pytorch_lightning/strategies           @tchaton @SeanNaren @awaelchli @justusschock @kaushikb11
 /src/pytorch_lightning/trainer              @williamfalcon @borda @tchaton @SeanNaren @carmocca @awaelchli @justusschock @kaushikb11
 /src/pytorch_lightning/trainer/connectors   @tchaton @SeanNaren @carmocca @borda

dockers/tpu-tests/tpu_test_cases.jsonnet

@@ -36,7 +36,7 @@ local tputests = base.BaseTest {
       # TODO (@kaushikb11): Add device stats tests here
       coverage run --source=pytorch_lightning -m pytest -v --capture=no \
           strategies/test_tpu_spawn.py \
-          profiler/test_xla_profiler.py \
+          profilers/test_xla_profiler.py \
           accelerators/test_tpu.py \
           models/test_tpu.py
       test_exit_code=$?

docs/source-pytorch/api_references.rst

@@ -231,7 +231,7 @@ others
 profiler
 --------
 
-.. currentmodule:: pytorch_lightning.profiler
+.. currentmodule:: pytorch_lightning.profilers
 
 .. autosummary::
     :toctree: api

docs/source-pytorch/common/trainer.rst

@@ -1213,7 +1213,7 @@ See the :doc:`profiler documentation <../tuning/profiler>`. for more details.
 
 .. testcode::
 
-    from pytorch_lightning.profiler import SimpleProfiler, AdvancedProfiler
+    from pytorch_lightning.profilers import SimpleProfiler, AdvancedProfiler
 
     # default used by the Trainer
     trainer = Trainer(profiler=None)

docs/source-pytorch/tuning/profiler_advanced.rst

@@ -12,11 +12,11 @@ Find bottlenecks in your code (advanced)
 ************************
 Profile cloud TPU models
 ************************
-To profile TPU models use the :class:`~pytorch_lightning.profiler.xla.XLAProfiler`
+To profile TPU models use the :class:`~pytorch_lightning.profilers.xla.XLAProfiler`
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import XLAProfiler
+    from pytorch_lightning.profilers import XLAProfiler
 
     profiler = XLAProfiler(port=9001)
     trainer = Trainer(profiler=profiler)

docs/source-pytorch/tuning/profiler_basic.rst

@@ -68,7 +68,7 @@ The simple profiler measures all the standard methods used in the training loop
 **************************************
 Profile the time within every function
 **************************************
-To profile the time within every function, use the :class:`~pytorch_lightning.profiler.advanced.AdvancedProfiler` built on top of Python's `cProfiler <https://docs.python.org/3/library/profile.html#module-cProfile>`_.
+To profile the time within every function, use the :class:`~pytorch_lightning.profilers.advanced.AdvancedProfiler` built on top of Python's `cProfiler <https://docs.python.org/3/library/profile.html#module-cProfile>`_.
 
 
 .. code-block:: python
@@ -101,7 +101,7 @@ If the profiler report becomes too long, you can stream the report to a file:
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import AdvancedProfiler
+    from pytorch_lightning.profilers import AdvancedProfiler
 
     profiler = AdvancedProfiler(dirpath=".", filename="perf_logs")
     trainer = Trainer(profiler=profiler)

docs/source-pytorch/tuning/profiler_expert.rst

@@ -12,12 +12,12 @@ Find bottlenecks in your code (expert)
 ***********************
 Build your own profiler
 ***********************
-To build your own profiler, subclass :class:`~pytorch_lightning.profiler.base.Profiler`
+To build your own profiler, subclass :class:`~pytorch_lightning.profilers.profiler.Profiler`
 and override some of its methods. Here is a simple example that profiles the first occurrence and total calls of each action:
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import Profiler
+    from pytorch_lightning.profilers import Profiler
     from collections import defaultdict
     import time
 
@@ -69,7 +69,7 @@ To profile a specific action of interest, reference a profiler in the LightningM
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import SimpleProfiler, PassThroughProfiler
+    from pytorch_lightning.profilers import SimpleProfiler, PassThroughProfiler
 
 
     class MyModel(LightningModule):
@@ -90,7 +90,7 @@ Here's the full code:
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import SimpleProfiler, PassThroughProfiler
+    from pytorch_lightning.profilers import SimpleProfiler, PassThroughProfiler
 
 
     class MyModel(LightningModule):

docs/source-pytorch/tuning/profiler_intermediate.rst

@@ -12,11 +12,11 @@ Find bottlenecks in your code (intermediate)
 **************************
 Profile pytorch operations
 **************************
-To understand the cost of each PyTorch operation, use the :class:`~pytorch_lightning.profiler.pytorch.PyTorchProfiler` built on top of the `PyTorch profiler <https://pytorch.org/docs/master/profiler.html>`__.
+To understand the cost of each PyTorch operation, use the :class:`~pytorch_lightning.profilers.pytorch.PyTorchProfiler` built on top of the `PyTorch profiler <https://pytorch.org/docs/master/profiler.html>`__.
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import PyTorchProfiler
+    from pytorch_lightning.profilers import PyTorchProfiler
 
     profiler = PyTorchProfiler()
     trainer = Trainer(profiler=profiler)
@@ -65,11 +65,11 @@ The profiler will generate an output like this:
 ***************************
 Profile a distributed model
 ***************************
-To profile a distributed model, use the :class:`~pytorch_lightning.profiler.pytorch.PyTorchProfiler` with the *filename* argument which will save a report per rank.
+To profile a distributed model, use the :class:`~pytorch_lightning.profilers.pytorch.PyTorchProfiler` with the *filename* argument which will save a report per rank.
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import PyTorchProfiler
+    from pytorch_lightning.profilers import PyTorchProfiler
 
     profiler = PyTorchProfiler(filename="perf-logs")
     trainer = Trainer(profiler=profiler)
@@ -153,11 +153,11 @@ to extend the scope of profiled functions.
 *****************************
 Visualize profiled operations
 *****************************
-To visualize the profiled operations, enable **emit_nvtx** in the :class:`~pytorch_lightning.profiler.pytorch.PyTorchProfiler`.
+To visualize the profiled operations, enable **emit_nvtx** in the :class:`~pytorch_lightning.profilers.pytorch.PyTorchProfiler`.
 
 .. code-block:: python
 
-    from pytorch_lightning.profiler import PyTorchProfiler
+    from pytorch_lightning.profilers import PyTorchProfiler
 
     profiler = PyTorchProfiler(emit_nvtx=True)
     trainer = Trainer(profiler=profiler)

pyproject.toml

@@ -78,10 +78,10 @@ module = [
     "pytorch_lightning.strategies.single_tpu",
     "pytorch_lightning.strategies.tpu_spawn",
     "pytorch_lightning.strategies.strategy",
-    "pytorch_lightning.profiler.advanced",
-    "pytorch_lightning.profiler.base",
-    "pytorch_lightning.profiler.pytorch",
-    "pytorch_lightning.profiler.simple",
+    "pytorch_lightning.profilers.advanced",
+    "pytorch_lightning.profilers.base",
+    "pytorch_lightning.profilers.pytorch",
+    "pytorch_lightning.profilers.simple",
     "pytorch_lightning.trainer.callback_hook",
     "pytorch_lightning.trainer.connectors.callback_connector",
     "pytorch_lightning.trainer.connectors.data_connector",

src/pytorch_lightning/profilers/__init__.py

@@ -0,0 +1,28 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pytorch_lightning.profilers.advanced import AdvancedProfiler
+from pytorch_lightning.profilers.base import PassThroughProfiler
+from pytorch_lightning.profilers.profiler import Profiler
+from pytorch_lightning.profilers.pytorch import PyTorchProfiler
+from pytorch_lightning.profilers.simple import SimpleProfiler
+from pytorch_lightning.profilers.xla import XLAProfiler
+
+__all__ = [
+    "Profiler",
+    "AdvancedProfiler",
+    "PassThroughProfiler",
+    "PyTorchProfiler",
+    "SimpleProfiler",
+    "XLAProfiler",
+]

src/pytorch_lightning/profilers/advanced.py

@@ -0,0 +1,91 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Profiler to check if there are any bottlenecks in your code."""
+import cProfile
+import io
+import logging
+import pstats
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+from pytorch_lightning.profilers.profiler import Profiler
+
+log = logging.getLogger(__name__)
+
+
+class AdvancedProfiler(Profiler):
+    """This profiler uses Python's cProfiler to record more detailed information about time spent in each function
+    call recorded during a given action.
+
+    The output is quite verbose and you should only use this if you want very detailed reports.
+    """
+
+    def __init__(
+        self,
+        dirpath: Optional[Union[str, Path]] = None,
+        filename: Optional[str] = None,
+        line_count_restriction: float = 1.0,
+    ) -> None:
+        """
+        Args:
+            dirpath: Directory path for the ``filename``. If ``dirpath`` is ``None`` but ``filename`` is present, the
+                ``trainer.log_dir`` (from :class:`~pytorch_lightning.loggers.tensorboard.TensorBoardLogger`)
+                will be used.
+
+            filename: If present, filename where the profiler results will be saved instead of printing to stdout.
+                The ``.txt`` extension will be used automatically.
+
+            line_count_restriction: this can be used to limit the number of functions
+                reported for each action. either an integer (to select a count of lines),
+                or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines)
+
+        Raises:
+            ValueError:
+                If you attempt to stop recording an action which was never started.
+        """
+        super().__init__(dirpath=dirpath, filename=filename)
+        self.profiled_actions: Dict[str, cProfile.Profile] = {}
+        self.line_count_restriction = line_count_restriction
+
+    def start(self, action_name: str) -> None:
+        if action_name not in self.profiled_actions:
+            self.profiled_actions[action_name] = cProfile.Profile()
+        self.profiled_actions[action_name].enable()
+
+    def stop(self, action_name: str) -> None:
+        pr = self.profiled_actions.get(action_name)
+        if pr is None:
+            raise ValueError(f"Attempting to stop recording an action ({action_name}) which was never started.")
+        pr.disable()
+
+    def summary(self) -> str:
+        recorded_stats = {}
+        for action_name, pr in self.profiled_actions.items():
+            s = io.StringIO()
+            ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats("cumulative")
+            ps.print_stats(self.line_count_restriction)
+            recorded_stats[action_name] = s.getvalue()
+        return self._stats_to_str(recorded_stats)
+
+    def teardown(self, stage: Optional[str] = None) -> None:
+        super().teardown(stage=stage)
+        self.profiled_actions = {}
+
+    def __reduce__(self):
+        # avoids `TypeError: cannot pickle 'cProfile.Profile' object`
+        return (
+            self.__class__,
+            (),
+            dict(dirpath=self.dirpath, filename=self.filename, line_count_restriction=self.line_count_restriction),
+        )

src/pytorch_lightning/profilers/base.py

@@ -0,0 +1,29 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Profiler to check if there are any bottlenecks in your code."""
+
+from pytorch_lightning.profilers.profiler import Profiler
+
+
+class PassThroughProfiler(Profiler):
+    """This class should be used when you don't want the (small) overhead of profiling.
+
+    The Trainer uses this class by default.
+    """
+
+    def start(self, action_name: str) -> None:
+        pass
+
+    def stop(self, action_name: str) -> None:
+        pass