Use latest pytorch-lightning version (#926)

Previously, we had fixed the version to an earlier release, which became incompatible after pip updated to 24.1. So, we took this opportunity to update it to the latest version.

dev-requirements.txt

@@ -20,7 +20,8 @@ pyre-extensions
 pyre-check
 pytest
 pytest-cov
-pytorch-lightning==1.5.10
+pytorch-lightning==2.3.1
+tensorboard==2.14.0
 sagemaker>=2.149.0
 torch-model-archiver>=0.4.2
 torch==2.2.1

torchx/examples/apps/lightning/profiler.py

@@ -19,18 +19,19 @@
 import time
 from typing import Dict
 
-from pytorch_lightning.loggers.base import LightningLoggerBase
-from pytorch_lightning.profiler.base import BaseProfiler
+from pytorch_lightning.loggers.logger import Logger
 
+from pytorch_lightning.profilers.profiler import Profiler
 
-class SimpleLoggingProfiler(BaseProfiler):
+
+class SimpleLoggingProfiler(Profiler):
     """
     This profiler records the duration of actions (in seconds) and reports the
     mean duration of each action to the specified logger. Reported metrics are
     in the format `duration_<event>`.
     """
 
-    def __init__(self, logger: LightningLoggerBase) -> None:
+    def __init__(self, logger: Logger) -> None:
         super().__init__()
 
         self.current_actions: Dict[str, float] = {}