-
Notifications
You must be signed in to change notification settings - Fork 3.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Enable dumping raw prof files in AdvancedProfiler
#19703
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -16,12 +16,15 @@ | |
import cProfile | ||
import io | ||
import logging | ||
import os | ||
import pstats | ||
import tempfile | ||
from pathlib import Path | ||
from typing import Dict, Optional, Tuple, Union | ||
|
||
from typing_extensions import override | ||
|
||
from lightning.fabric.utilities.cloud_io import get_filesystem | ||
from lightning.pytorch.profilers.profiler import Profiler | ||
|
||
log = logging.getLogger(__name__) | ||
|
@@ -40,6 +43,7 @@ def __init__( | |
dirpath: Optional[Union[str, Path]] = None, | ||
filename: Optional[str] = None, | ||
line_count_restriction: float = 1.0, | ||
dump_stats: bool = False, | ||
) -> None: | ||
""" | ||
Args: | ||
|
@@ -54,13 +58,16 @@ def __init__( | |
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) | ||
|
||
dump_stats: Whether to save raw profiler results. When ``True`` then ``dirpath`` must be provided. | ||
|
||
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 | ||
self.dump_stats = dump_stats | ||
|
||
@override | ||
def start(self, action_name: str) -> None: | ||
|
@@ -75,10 +82,27 @@ def stop(self, action_name: str) -> None: | |
raise ValueError(f"Attempting to stop recording an action ({action_name}) which was never started.") | ||
pr.disable() | ||
|
||
def _dump_stats(self, action_name: str, profile: cProfile.Profile) -> None: | ||
assert self.dirpath | ||
dst_filepath = os.path.join(self.dirpath, self._prepare_filename(action_name=action_name, extension=".prof")) | ||
dst_fs = get_filesystem(dst_filepath) | ||
dst_fs.mkdirs(self.dirpath, exist_ok=True) | ||
# temporarily save to local since pstats can only dump into a local file | ||
with tempfile.TemporaryDirectory(prefix="test", suffix="test", dir=os.getcwd()) as tmp_dir, dst_fs.open( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we'd need to include the rank in the prefix. Otherwise we could run into rare race conditions, right? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not really. The temp directory will be unique for each rank so there will be no collision in the full file path. The directory is removed right after this contextmanager exits. The file name already adds the rank in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. For the filename that goes into that directory, yes. But the temp directory itself
is just a randomized name saved in the current working directory. There is no guarantee that the name of this directory will be different on each rank. |
||
dst_filepath, "wb" | ||
) as dst_file: | ||
src_filepath = os.path.join(tmp_dir, "tmp.prof") | ||
profile.dump_stats(src_filepath) | ||
src_fs = get_filesystem(src_filepath) | ||
with src_fs.open(src_filepath, "rb") as src_file: | ||
dst_file.write(src_file.read()) | ||
|
||
@override | ||
def summary(self) -> str: | ||
recorded_stats = {} | ||
for action_name, pr in self.profiled_actions.items(): | ||
if self.dump_stats: | ||
self._dump_stats(action_name, pr) | ||
s = io.StringIO() | ||
ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats("cumulative") | ||
ps.print_stats(self.line_count_restriction) | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -299,6 +299,25 @@ def test_advanced_profiler_describe(tmp_path, advanced_profiler): | |
assert len(data) > 0 | ||
|
||
|
||
def test_advanced_profiler_dump_states(tmp_path): | ||
advanced_profiler = AdvancedProfiler(dirpath=tmp_path, dump_stats=True) | ||
"""Ensure the profiler dump stats during summary.""" | ||
# record at least one event | ||
with advanced_profiler.profile(action_name := "test"): | ||
pass | ||
# dump_stats to file | ||
advanced_profiler.describe() | ||
path = advanced_profiler.dirpath / f"{action_name}.prof" | ||
data = path.read_bytes() | ||
assert len(data) > 0 | ||
|
||
|
||
def test_advanced_profiler_dump_states_needs_dirpath(): | ||
"""Ensure the profiler requires dirpath to dump stats.""" | ||
with pytest.raises(AssertionError): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This doesn't seem to be raised anymore. Let's remove it. |
||
AdvancedProfiler(dump_stats=True) | ||
|
||
|
||
def test_advanced_profiler_value_errors(advanced_profiler): | ||
"""Ensure errors are raised where expected.""" | ||
action = "test" | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.