Skip to content

Commit

Permalink
Fix action_name usage in XLAProfiler (#15886)
Browse files Browse the repository at this point in the history
* Fix `action_name` usage in `XLAProfiler`
* add changelog
* Update src/pytorch_ligh
* Update xla.py

Co-authored-by: awaelchli <aedu.waelchli@gmail.com>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
  • Loading branch information
3 people authored Dec 9, 2022
1 parent a72d268 commit c748f82
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
3 changes: 3 additions & 0 deletions src/pytorch_lightning/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Enhanced `reduce_boolean_decision` to accommodate `any`-analogous semantics expected by the `EarlyStopping` callback ([#15253](https://github.com/Lightning-AI/lightning/pull/15253))


- Fixed the `XLAProfiler` not recording anything due to mismatching of action names ([#15885](https://github.com/Lightning-AI/lightning/pull/15885))


## [1.8.4] - 2022-12-08

### Changed
Expand Down
6 changes: 4 additions & 2 deletions src/pytorch_lightning/profilers/xla.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,12 +50,14 @@ def __init__(self, port: int = 9012) -> None:
def start(self, action_name: str) -> None:
import torch_xla.debug.profiler as xp

if action_name in self.RECORD_FUNCTIONS:
# The action name is formatted as '[TYPE]{class name}.{hook name}'
# Example: [LightningModule]BoringModel.training_step
if action_name.split(".")[-1] in self.RECORD_FUNCTIONS:
if not self._start_trace:
self.server = xp.start_server(self.port)
self._start_trace = True

if action_name in self.STEP_FUNCTIONS:
if action_name.split(".")[-1] in self.STEP_FUNCTIONS:
step = self._get_step_num(action_name)
recording = xp.StepTrace(action_name, step_num=step)
else:
Expand Down

0 comments on commit c748f82

Please sign in to comment.