Fix mypy for utilities.xla_device (Lightning-AI#8755)

* Fix mypy for utilities.xla_device

* Add explicit type hint for _XLA_AVAILABLE in utilities.imports

Author: stancld

pyproject.toml

@@ -71,5 +71,6 @@ module = [
     "pytorch_lightning.utilities.distributed",
     "pytorch_lightning.utilities.memory",
     "pytorch_lightning.utilities.parsing",
+    "pytorch_lightning.utilities.xla_device",
 ]
 ignore_errors = "False"

pytorch_lightning/utilities/imports.py

@@ -89,7 +89,7 @@ def _compare_version(package: str, op, version) -> bool:
 _TORCHVISION_AVAILABLE = _module_available("torchvision")
 _TORCHMETRICS_LOWER_THAN_0_3 = _compare_version("torchmetrics", operator.lt, "0.3.0")
 _TORCHMETRICS_GREATER_EQUAL_0_3 = _compare_version("torchmetrics", operator.ge, "0.3.0")
-_XLA_AVAILABLE = _module_available("torch_xla")
+_XLA_AVAILABLE: bool = _module_available("torch_xla")
 
 from pytorch_lightning.utilities.xla_device import XLADeviceUtils  # noqa: E402
 

pytorch_lightning/utilities/xla_device.py

@@ -16,6 +16,7 @@
 import queue as q
 import traceback
 from multiprocessing import Process, Queue
+from typing import Any, Callable, Union
 
 from pytorch_lightning.utilities.imports import _XLA_AVAILABLE
 
@@ -26,7 +27,7 @@
 TPU_CHECK_TIMEOUT = 60
 
 
-def inner_f(queue, func, *args, **kwargs):  # pragma: no cover
+def inner_f(queue: Queue, func: Callable, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
     try:
         queue.put(func(*args, **kwargs))
     # todo: specify the possible exception
@@ -35,10 +36,10 @@ def inner_f(queue, func, *args, **kwargs):  # pragma: no cover
         queue.put(None)
 
 
-def pl_multi_process(func):
+def pl_multi_process(func: Callable) -> Callable:
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        queue = Queue()
+    def wrapper(*args: Any, **kwargs: Any) -> Union[bool, Any]:
+        queue: Queue = Queue()
         proc = Process(target=inner_f, args=(queue, func, *args), kwargs=kwargs)
         proc.start()
         proc.join(TPU_CHECK_TIMEOUT)