Skip to content

Commit

Permalink
Reroute profiler to profilers (Lightning-AI#12308)
Browse files Browse the repository at this point in the history 
Co-authored-by: Carlos Mocholí <[email protected]>
Co-authored-by: Akihiro Nitta <[email protected]>
  • Loading branch information
@awaelchli @carmocca @akihironitta
3 people committed Jun 21, 2022
1 parent 29d6cda commit 49a4ed7
Show file tree
Hide file tree
Showing 14 changed files with 131 additions and 940 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Deprecated LightningCLI's registries in favor of importing the respective package ([#13221](https://github.com/PyTorchLightning/pytorch-lightning/pull/13221))



- Deprecated `pytorch_lightning.profiler` in favor of `pytorch_lightning.profilers` ([#12308](https://github.com/PyTorchLightning/pytorch-lightning/pull/12308))


-

### Removed

- Removed the deprecated `Logger.close` method ([#13149](https://github.com/PyTorchLightning/pytorch-lightning/pull/13149))
Expand Down
13 changes: 7 additions & 6 deletions src/pytorch_lightning/profiler/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,13 @@
# 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.profiler.advanced import AdvancedProfiler
from pytorch_lightning.profiler.base import AbstractProfiler, BaseProfiler, PassThroughProfiler
from pytorch_lightning.profiler.profiler import Profiler
from pytorch_lightning.profiler.pytorch import PyTorchProfiler
from pytorch_lightning.profiler.simple import SimpleProfiler
from pytorch_lightning.profiler.xla import XLAProfiler
from pytorch_lightning.profiler.base import AbstractProfiler, BaseProfiler
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__ = [
"AbstractProfiler",
Expand Down
83 changes: 8 additions & 75 deletions src/pytorch_lightning/profiler/advanced.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,81 +11,14 @@
# 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.advanced import AdvancedProfiler as NewAdvancedProfiler
from pytorch_lightning.utilities import rank_zero_deprecation

from pytorch_lightning.profiler.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),
class AdvancedProfiler(NewAdvancedProfiler):
def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def]
rank_zero_deprecation(
"`pytorch_lightning.profiler.AdvancedProfiler` is deprecated in v1.7 and will be removed in v1.9."
" Use the equivalent `pytorch_lightning.profilers.AdvancedProfiler` class instead."
)
super().__init__(*args, **kwargs)
23 changes: 10 additions & 13 deletions src/pytorch_lightning/profiler/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
from abc import ABC, abstractmethod
from typing import Any

from pytorch_lightning.profiler.profiler import Profiler
from pytorch_lightning.profilers.base import PassThroughProfiler as NewPassThroughProfiler
from pytorch_lightning.profilers.profiler import Profiler
from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation


Expand Down Expand Up @@ -57,21 +58,17 @@ class BaseProfiler(Profiler):
Please use `Profiler` instead.
"""

def __init__(self, *args, **kwargs):
def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def]
rank_zero_deprecation(
"`BaseProfiler` was deprecated in v1.6 and will be removed in v1.8. Please use `Profiler` instead."
)
super().__init__(*args, **kwargs)


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
class PassThroughProfiler(NewPassThroughProfiler):
def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def]
rank_zero_deprecation(
"`pytorch_lightning.profiler.PassThroughProfiler` is deprecated in v1.7 and will be removed in v1.9."
" Use the equivalent `pytorch_lightning.profilers.PassThroughProfiler` class instead."
)
super().__init__(*args, **kwargs)
168 changes: 12 additions & 156 deletions src/pytorch_lightning/profiler/profiler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,164 +11,20 @@
# 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 logging
import os
from abc import ABC, abstractmethod
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Callable, Dict, Generator, Iterable, Optional, TextIO, Union
from pytorch_lightning.profilers.profiler import Profiler as NewProfiler
from pytorch_lightning.utilities import rank_zero_deprecation

from pytorch_lightning.utilities.cloud_io import get_filesystem
from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation

log = logging.getLogger(__name__)
class Profiler(NewProfiler):
"""
.. deprecated:: v1.6
`pytorch_lightning.profiler.Profiler` is deprecated in v1.7 and will be removed in v1.9.
Use the equivalent `pytorch_lightning.profilers.Profiler` class instead.
"""


class Profiler(ABC):
"""If you wish to write a custom profiler, you should inherit from this class."""

def __init__(
self,
dirpath: Optional[Union[str, Path]] = None,
filename: Optional[str] = None,
) -> None:
self.dirpath = dirpath
self.filename = filename

self._output_file: Optional[TextIO] = None
self._write_stream: Optional[Callable] = None
self._local_rank: Optional[int] = None
self._stage: Optional[str] = None

@abstractmethod
def start(self, action_name: str) -> None:
"""Defines how to start recording an action."""

@abstractmethod
def stop(self, action_name: str) -> None:
"""Defines how to record the duration once an action is complete."""

def summary(self) -> str:
return ""

@contextmanager
def profile(self, action_name: str) -> Generator:
"""Yields a context manager to encapsulate the scope of a profiled action.
Example::
with self.profile('load training data'):
# load training data code
The profiler will start once you've entered the context and will automatically
stop once you exit the code block.
"""
try:
self.start(action_name)
yield action_name
finally:
self.stop(action_name)

def profile_iterable(self, iterable: Iterable, action_name: str) -> Generator:
"""Profiles over each value of an iterable.
See deprecation message below.
.. deprecated:: v1.6
`Profiler.profile_iterable` is deprecated in v1.6 and will be removed in v1.8.
"""
def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def]
rank_zero_deprecation(
f"`{self.__class__.__name__}.profile_iterable` is deprecated in v1.6 and will be removed in v1.8."
"`pytorch_lightning.profiler.Profiler` is deprecated in v1.7 and will be removed in v1.9."
" Use the equivalent `pytorch_lightning.profilers.Profiler` class instead."
)
iterator = iter(iterable)
while True:
try:
self.start(action_name)
value = next(iterator)
self.stop(action_name)
yield value
except StopIteration:
self.stop(action_name)
break

def _rank_zero_info(self, *args: Any, **kwargs: Any) -> None:
if self._local_rank in (None, 0):
log.info(*args, **kwargs)

def _prepare_filename(
self, action_name: Optional[str] = None, extension: str = ".txt", split_token: str = "-"
) -> str:
args = []
if self._stage is not None:
args.append(self._stage)
if self.filename:
args.append(self.filename)
if self._local_rank is not None:
args.append(str(self._local_rank))
if action_name is not None:
args.append(action_name)
filename = split_token.join(args) + extension
return filename

def _prepare_streams(self) -> None:
if self._write_stream is not None:
return
if self.filename and self.dirpath:
filepath = os.path.join(self.dirpath, self._prepare_filename())
fs = get_filesystem(filepath)
fs.mkdirs(self.dirpath, exist_ok=True)
file = fs.open(filepath, "a")
self._output_file = file
self._write_stream = file.write
else:
self._write_stream = self._rank_zero_info

def describe(self) -> None:
"""Logs a profile report after the conclusion of run."""
# users might call `describe` directly as the profilers can be used by themselves.
# to allow this, we open and close the files within this function by calling `_prepare_streams` and `teardown`
# manually instead of letting the `Trainer` do it through `setup` and `teardown`
self._prepare_streams()
summary = self.summary()
if summary and self._write_stream is not None:
self._write_stream(summary)
if self._output_file is not None:
self._output_file.flush()
self.teardown(stage=self._stage)

def _stats_to_str(self, stats: Dict[str, str]) -> str:
stage = f"{self._stage.upper()} " if self._stage is not None else ""
output = [stage + "Profiler Report"]
for action, value in stats.items():
header = f"Profile stats for: {action}"
if self._local_rank is not None:
header += f" rank: {self._local_rank}"
output.append(header)
output.append(value)
return os.linesep.join(output)

def setup(
self, stage: Optional[str] = None, local_rank: Optional[int] = None, log_dir: Optional[str] = None
) -> None:
"""Execute arbitrary pre-profiling set-up steps."""
self._stage = stage
self._local_rank = local_rank
self.dirpath = self.dirpath or log_dir

def teardown(self, stage: Optional[str] = None) -> None:
"""Execute arbitrary post-profiling tear-down steps.
Closes the currently open file and stream.
"""
self._write_stream = None
if self._output_file is not None:
self._output_file.close()
self._output_file = None # can't pickle TextIOWrapper

def __del__(self) -> None:
self.teardown(stage=self._stage)

@property
def local_rank(self) -> int:
return 0 if self._local_rank is None else self._local_rank
super().__init__(*args, **kwargs)
Loading

0 comments on commit 49a4ed7

Please sign in to comment.