Repository: Lightning-AI/pytorch-lightning
Pull request: Fix `action_name` usage in `XLAProfiler`
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #15885

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
`XLAProfiler` on TPU not working
### Bug description

It seems the changes in `action_name` is preventing it from starting the profiler server or logging any steps. This patch fixes the bug:


```
diff --git a/src/pytorch_lightning/profilers/xla.py b/src/pytorch_lightning/profilers/xla.py
index ef103a9a4..958e8de3f 100644
--- a/src/pytorch_lightning/profilers/xla.py
+++ b/src/pytorch_lightning/profilers/xla.py
@@ -50,7 +50,7 @@ class XLAProfiler(Profiler):
     def start(self, action_name: str) -> None:
         import torch_xla.debug.profiler as xp
 
-        if action_name in self.RECORD_FUNCTIONS:
+        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
```



### How to reproduce the bug

_No response_

### Error messages and logs

```

# Error messages and logs here please

```


### Environment

```

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 1.10):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

```


### More info

_No response_

cc @carmocca @nbcsm @guotuofeng

Hints from issue discussion:
Hey @@Liyang90,

Would you be interested into turning this patch into a PR ?

Best,
T.C

Commit message: Fix `action_name` usage in `XLAProfiler`

Patch to review:
diff --git a/src/pytorch_lightning/profilers/xla.py b/src/pytorch_lightning/profilers/xla.py
index ef103a9a45842..958e8de3f7e18 100644
--- a/src/pytorch_lightning/profilers/xla.py
+++ b/src/pytorch_lightning/profilers/xla.py
@@ -50,7 +50,7 @@ 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:
+        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


Code review comments:
diff_hunk: @@ -50,7 +50,7 @@ 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:
+        if action_name.split(".")[-1] in self.RECORD_FUNCTIONS:
path: src/pytorch_lightning/profilers/xla.py
text: @user1:
Don't we need to apply the same fix on line 58?

@user1:
Additionally, I suggest adding a comment here:

```suggestion
        # 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:
```