class_docstr: Launches processes that run a given function in parallel on XLA supported hardware, and joins them all at the
  end.
  
  The main process in which this launcher is invoked creates N so-called worker processes (using the
  `torch_xla` :func:`xmp.spawn`) that run the given function.
  Worker processes have a rank that ranges from 0 to N - 1.
  
  Note:
      - This launcher requires all objects to be pickleable.
      - It is important that the entry point to the program/script is guarded by ``if __name__ == "__main__"``.
  
  Args:
      strategy: A reference to the strategy that is used together with this launcher
class_docstr_tok: ["Launches","processes","that","run","a","given","function","in","parallel","on","XLA","supported","hardware",",","and","joins","them","all","at","the","end",".","The","main","process","in","which","this","launcher","is","invoked","creates","N","so-called","worker","processes","(","using","the","`","torch_xla","`",":","func",":","`","xmp.spawn","`",")","that","run","the","given","function",".","Worker","processes","have","a","rank","that","ranges","from","0","to","N","-","1",".","Note",":","-","This","launcher","requires","all","objects","to","be","pickleable",".","-","It","is","important","that","the","entry","point","to","the","program\/script","is","guarded","by","``","if","__name__","==","``","__main__","''","``",".","Args",":","strategy",":","A","reference","to","the","strategy","that","is","used","together","with","this","launcher"]
docstr: Launches processes that run the given function in parallel.
  
  The function is allowed to have a return value. However, when all processes join, only the return value
  of worker process 0 gets returned from this `launch` method in the main process.
  
  Arguments:
      function: The entry point for all launched processes.
      *args: Optional positional arguments to be passed to the given function.
      trainer: Optional reference to the :class:`~lightning.pytorch.trainer.trainer.Trainer` for which
          a selected set of attributes get restored in the main process after processes join.
      **kwargs: Optional keyword arguments to be passed to the given function.
docstr_tok: ["Launches","processes","that","run","the","given","function","in","parallel",".","The","function","is","allowed","to","have","a","return","value",".","However",",","when","all","processes","join",",","only","the","return","value","of","worker","process","0","gets","returned","from","this","`","launch","`","method","in","the","main","process",".","Arguments",":","function",":","The","entry","point","for","all","launched","processes",".","*","args",":","Optional","positional","arguments","to","be","passed","to","the","given","function",".","trainer",":","Optional","reference","to","the",":","class",":","`","~lightning.pytorch.trainer.trainer.Trainer","`","for","which","a","selected","set","of","attributes","get","restored","in","the","main","process","after","processes","join",".","*","*","kwargs",":","Optional","keyword","arguments","to","be","passed","to","the","given","function","."]
code: def launch(
      self,
      function: Callable,
      *args: Any,
      trainer: Optional["pl.Trainer"] = None,
      **kwargs: Any,
  ) -> Any:
      """Launches processes that run the given function in parallel.
  
      The function is allowed to have a return value. However, when all processes join, only the return value
      of worker process 0 gets returned from this `launch` method in the main process.
  
      Arguments:
          function: The entry point for all launched processes.
          *args: Optional positional arguments to be passed to the given function.
          trainer: Optional reference to the :class:`~lightning.pytorch.trainer.trainer.Trainer` for which
              a selected set of attributes get restored in the main process after processes join.
          **kwargs: Optional keyword arguments to be passed to the given function.
  
      """
      if (
          self._already_fit
          and trainer is not None
          and trainer.state.fn == TrainerFn.FITTING
      ):
          # resolving https://github.com/Lightning-AI/lightning/issues/18775 will lift this restriction
          raise NotImplementedError(
              "Calling `trainer.fit()` twice on the same Trainer instance using a spawn-based strategy is not"
              " supported. You can work around this by creating a new Trainer instance and passing the"
              " `fit(ckpt_path=...)` argument."
          )
  
      # pjrt requires that the queue is serializable
      return_queue = mp.Manager().Queue()
  
      import torch_xla.distributed.xla_multiprocessing as xmp
  
      spawn_kwargs = {}
      nprocs = self._strategy.num_processes
      if nprocs == 1:
          # avoid warning: "Unsupported nprocs". If it's 1, it will call the launched function directly.
          # otherwise it will use all devices
          spawn_kwargs["nprocs"] = nprocs
  
      process_context = xmp.spawn(
          self._wrapping_function,
          args=(trainer, function, args, kwargs, return_queue),
          start_method=self._start_method,
          join=False,  # we will join ourselves to get the process references
          **spawn_kwargs,
      )
      # xla will not actually create processes if only 1 device
      if process_context is not None:
          self.procs = process_context.processes
          while not process_context.join():
              pass
  
      worker_output = return_queue.get()
      if trainer is None:
          return worker_output
  
      self._already_fit |= trainer.state.fn == TrainerFn.FITTING
      self._recover_results_in_main_process(worker_output, trainer)
      return worker_output.trainer_results
  
file_path: pytorch-lightning/src/lightning/pytorch/strategies/launchers/xla.py
imports: import os
  import queue
  from typing import TYPE_CHECKING, Any, Callable, Optional, Union
  import torch.multiprocessing
  from typing_extensions import override
  from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
  from lightning.fabric.strategies.launchers.xla import _rank_teardown
  from lightning.fabric.utilities import move_data_to_device
  from lightning.pytorch.strategies.launchers.multiprocessing import _GlobalStateSnapshot, _MultiProcessingLauncher, _WorkerOutput
  from lightning.pytorch.trainer.states import TrainerFn
  from lightning.pytorch.utilities.rank_zero import rank_zero_debug
label_desc: Use image node_id 3 for calling the _XLALauncher obj's underlying member method code with example usage: obj.launch(function) and returns: worker_output, worker_output