import torch.mp simplequeue class

Author: matsumotosan

src/lightning/pytorch/strategies/launchers/multiprocessing.py

@@ -26,6 +26,7 @@
 import torch.multiprocessing as mp
 from lightning_utilities.core.apply_func import apply_to_collection
 from torch import Tensor
+from torch.multiprocessing.queue import SimpleQueue
 from typing_extensions import override
 
 import lightning.pytorch as pl
@@ -160,7 +161,7 @@ def _wrapping_function(
         function: Callable,
         args: Any,
         kwargs: Any,
-        return_queue: mp.SimpleQueue | queue.Queue,
+        return_queue: SimpleQueue | queue.Queue,
         global_states: Optional["_GlobalStateSnapshot"] = None,
     ) -> None:
         if global_states:

src/lightning/pytorch/strategies/launchers/xla.py

@@ -17,6 +17,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch.multiprocessing as mp
+from torch.multiprocessing.queue import SimpleQueue
 from typing_extensions import override
 
 from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
@@ -127,7 +128,7 @@ def _wrapping_function(
         function: Callable,
         args: Any,
         kwargs: Any,
-        return_queue: mp.SimpleQueue | queue.Queue,
+        return_queue: SimpleQueue | queue.Queue,
         global_states: _GlobalStateSnapshot | None = None,
     ) -> None:
         import torch_xla.core.xla_model as xm