Propagate spawn kwargs from parallel to model's spawn

ignite/distributed/comp_models/xla.py

@@ -113,9 +113,7 @@ def spawn(
             backend: str = XLA_TPU,
             **kwargs
         ):
-            import os
-
-            if "COLAB_TPU_ADDR" in os.environ:
+            if "start_method" not in kwargs:
                 kwargs["start_method"] = "fork"
 
             xmp.spawn(

ignite/distributed/launcher.py

@@ -150,6 +150,7 @@ def training(local_rank, config, **kwargs):
             (`nccl`, `gloo`). Mandatory argument if ``nnodes`` is specified and larger than one.
         master_port (int, optional): optional argument, master node port for torch native backends
             (`nccl`, `gloo`). Mandatory argument if ``master_addr`` is specified.
+        **spawn_kwargs: kwargs to ``idist.spawn`` function.
     """
 
     def __init__(
@@ -160,6 +161,7 @@ def __init__(
         node_rank: Optional[int] = None,
         master_addr: Optional[str] = None,
         master_port: Optional[str] = None,
+        **spawn_kwargs
     ):
         if backend is not None:
             if backend not in idist.available_backends():
@@ -183,15 +185,16 @@ def __init__(
         if self.backend is not None:
             if nproc_per_node is not None:
                 self._spawn_params = self._setup_spawn_params(
-                    nproc_per_node, nnodes, node_rank, master_addr, master_port
+                    nproc_per_node, nnodes, node_rank, master_addr, master_port, **spawn_kwargs
                 )
 
         if self._spawn_params is not None:
             self.logger.info("Initialized distributed launcher with backend: '{}'".format(self.backend))
             msg = "\n\t".join(["{}: {}".format(k, v) for k, v in self._spawn_params.items() if v is not None])
             self.logger.info("- Parameters to spawn processes: \n\t{}".format(msg))
 
-    def _setup_spawn_params(self, nproc_per_node, nnodes, node_rank, master_addr, master_port):
+    @staticmethod
+    def _setup_spawn_params(nproc_per_node, nnodes, node_rank, master_addr, master_port, **spawn_kwargs):
         if nproc_per_node < 1:
             raise ValueError("Argument nproc_per_node should positive, but given {}".format(nproc_per_node))
         if nnodes is None:
@@ -218,6 +221,7 @@ def _setup_spawn_params(self, nproc_per_node, nnodes, node_rank, master_addr, ma
             "master_addr": master_addr,
             "master_port": master_port,
         }
+        params.update(spawn_kwargs)
         return {k: v for k, v in params.items() if v is not None}
 
     def run(self, func: Callable, *args, **kwargs):

tests/ignite/distributed/test_launcher.py

@@ -158,3 +158,21 @@ def test_idist_parallel_no_dist():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     with idist.Parallel(backend=None) as parallel:
         parallel.run(_test_func, ws=1, device=device)
+
+
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if no NUM_TPU_WORKERS in env vars")
+@pytest.mark.skipif(not has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_idist_parallel_spawn_params():
+
+    res = idist.Parallel._setup_spawn_params(
+        nproc_per_node=8, nnodes=None, node_rank=None, master_addr=None, master_port=None, start_method="fork"
+    )
+    assert "nproc_per_node" in res and res["nproc_per_node"] == 8
+    assert "start_method" in res and res["start_method"] == "fork"
+
+    with idist.Parallel(backend="xla-tpu", nproc_per_node=8, start_method="fork") as parallel:
+        assert parallel.backend == "xla-tpu"
+        res = parallel._spawn_params
+        assert "nproc_per_node" in res and res["nproc_per_node"] == 8
+        assert "start_method" in res and res["start_method"] == "fork"