XLAProfiler on TPU not working

[Liyang90]

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

[tchaton]

Hey @@Liyang90,

Would you be interested into turning this patch into a PR ?

Best,
T.C

[awaelchli]

@Liyang90 Thanks for reporting. Could you provide us with the error trace and the name of the action that caused the issue?

[Liyang90]

There is no error with in the Lightning, but the profiler would just not run because of the updates in the action name in the Trainer, such as: https://github.com/Lightning-AI/lightning/blob/12c74f134e4db63aa2f114fa539f1b45c721a1b8/src/pytorch_lightning/trainer/trainer.py#L1438
The name is not simply "training_step" or "backward" anymore.

[awaelchli]

I see. So either the dots or the brackets in the name seem to cause an issue. What if we do a string replacement in the xla profiler and replace these characters with another symbol? We could replace "." with "-" and maybe drop the brackets completely.

[Liyang90]

It doesn't seem to be caused by the special symbols. It the way the action names are filtered by self.RECORD_FUNCTIONS. It only has https://github.com/Lightning-AI/lightning/blob/12c74f134e4db63aa2f114fa539f1b45c721a1b8/src/pytorch_lightning/profilers/xla.py#L26.
But after some version, these only match the last part of the actual action names passed to the profiler. The linked PR fixes this.

[awaelchli]

@Liyang90 Oh I understand now. Sorry for my misleading comments. It's just that I haven't been active on that part of the code base very much.