fix type hints in XLA, SLURM, and Comet logger modules for improved clarity

Author: bhimrazy

src/lightning/fabric/accelerators/xla.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import functools
+import warnings
 from typing import Any
 
 import torch

src/lightning/fabric/plugins/environments/slurm.py

@@ -56,7 +56,7 @@ def auto_requeue(self) -> bool:
         return self.slurm_impl.auto_requeue
 
     @property
-    def requeue_signal(self) -> Optional[signal.Signals]:
+    def requeue_signal(self) -> signal.Signals | None:
         return self.slurm_impl.requeue_signal
 
     @property
@@ -93,7 +93,7 @@ def job_name() -> str | None:
         return os.environ.get("SLURM_JOB_NAME")
 
     @staticmethod
-    def job_id() -> Optional[int]:
+    def job_id() -> int | None:
         _raise_enterprise_not_available()
         from pytorch_lightning_enterprise.plugins.environments.slurm import (
             SLURMEnvironment as EnterpriseSLURMEnvironment,

src/lightning/pytorch/loggers/comet.py

@@ -242,7 +242,7 @@ def log_hyperparams(self, params: dict[str, Any] | Namespace) -> None:
 
     @override
     @rank_zero_only
-    def log_metrics(self, metrics: Mapping[str, Tensor | float], step: Optional[int] = None) -> None:
+    def log_metrics(self, metrics: Mapping[str, Tensor | float], step: int | None = None) -> None:
         return self.logger_impl.log_metrics(metrics, step)
 
     @override
@@ -287,5 +287,5 @@ def version(self) -> str | None:
         return self.logger_impl.version
 
     @override
-    def log_graph(self, model: Module, input_array: Optional[Tensor] = None) -> None:
+    def log_graph(self, model: Module, input_array: Tensor | None = None) -> None:
         return self.logger_impl.log_graph(model, input_array)