diff --git a/pytorch_lightning/plugins/training_type/deepspeed.py b/pytorch_lightning/plugins/training_type/deepspeed.py index f7223865869ed..fe3f51fa99390 100644 --- a/pytorch_lightning/plugins/training_type/deepspeed.py +++ b/pytorch_lightning/plugins/training_type/deepspeed.py @@ -30,7 +30,7 @@ from pytorch_lightning.trainer.optimizers import _get_default_scheduler_config from pytorch_lightning.utilities import AMPType from pytorch_lightning.utilities.apply_func import apply_to_collection -from pytorch_lightning.utilities.distributed import _warn, rank_zero_info, rank_zero_only +from pytorch_lightning.utilities.distributed import rank_zero_info, rank_zero_only from pytorch_lightning.utilities.exceptions import MisconfigurationException from pytorch_lightning.utilities.imports import _DEEPSPEED_AVAILABLE diff --git a/pytorch_lightning/profiler/profilers.py b/pytorch_lightning/profiler/profilers.py index 78327fa0a91d8..c91f271823e86 100644 --- a/pytorch_lightning/profiler/profilers.py +++ b/pytorch_lightning/profiler/profilers.py @@ -120,14 +120,19 @@ def _rank_zero_info(self, *args, **kwargs) -> None: if self._local_rank in (None, 0): log.info(*args, **kwargs) - def _prepare_filename(self, extension: str = ".txt") -> str: - filename = "" + def _prepare_filename( + self, action_name: Optional[str] = None, extension: str = ".txt", split_token: str = "-" + ) -> str: + args = [] if self._stage is not None: - filename += f"{self._stage}-" - filename += str(self.filename) + args.append(self._stage) + if self.filename: + args.append(self.filename) if self._local_rank is not None: - filename += f"-{self._local_rank}" - filename += extension + 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: diff --git a/pytorch_lightning/profiler/pytorch.py b/pytorch_lightning/profiler/pytorch.py index 9f78b194b56fc..ae59596d19e48 100644 --- a/pytorch_lightning/profiler/pytorch.py +++ b/pytorch_lightning/profiler/pytorch.py @@ -427,11 +427,15 @@ def stop(self, action_name: str) -> None: def on_trace_ready(profiler): if self.dirpath is not None: if self._export_to_chrome: - handler = tensorboard_trace_handler(self.dirpath, self._prepare_filename(extension="")) + handler = tensorboard_trace_handler( + self.dirpath, self._prepare_filename(action_name=action_name, extension="") + ) handler(profiler) if self._export_to_flame_graph: - path = os.path.join(self.dirpath, self._prepare_filename(extension=".stack")) + path = os.path.join( + self.dirpath, self._prepare_filename(action_name=action_name, extension=".stack") + ) profiler.export_stacks(path, metric=self._metric) else: rank_zero_warn("The PyTorchProfiler failed to export trace as `dirpath` is None") diff --git a/tests/deprecated_api/test_remove_1-5.py b/tests/deprecated_api/test_remove_1-5.py index bf6bb2bc3ba5a..22192b3d437f9 100644 --- a/tests/deprecated_api/test_remove_1-5.py +++ b/tests/deprecated_api/test_remove_1-5.py @@ -30,7 +30,7 @@ from pytorch_lightning.utilities import device_parser from pytorch_lightning.utilities.imports import _compare_version from tests.deprecated_api import no_deprecated_call -from tests.helpers import BoringDataModule, BoringModel +from tests.helpers import BoringModel from tests.helpers.utils import no_warning_call diff --git a/tests/test_profiler.py b/tests/test_profiler.py index 2a7620899ff39..f27365cee27f8 100644 --- a/tests/test_profiler.py +++ b/tests/test_profiler.py @@ -330,8 +330,8 @@ def test_pytorch_profiler_trainer_ddp(tmpdir, pytorch_profiler): files = [file for file in files if file.endswith('.json')] assert len(files) == 2, files local_rank = trainer.local_rank - assert any(f'training_step_{local_rank}' in f for f in files) - assert any(f'validation_step_{local_rank}' in f for f in files) + assert any(f'{local_rank}-training_step_and_backward' in f for f in files) + assert any(f'{local_rank}-validation_step' in f for f in files) def test_pytorch_profiler_trainer_test(tmpdir):