[App] Improve LightningTrainerScript start-up time

examples/app_multi_node/train_pytorch.py

@@ -23,7 +23,7 @@ def distributed_train(local_rank: int, main_address: str, main_port: int, num_no
     # 2. PREPARE DISTRIBUTED MODEL
     model = torch.nn.Linear(32, 2)
     device = torch.device(f"cuda:{local_rank}") if torch.cuda.is_available() else torch.device("cpu")
-    model = DistributedDataParallel(model, device_ids=[local_rank]).to(device)
+    model = DistributedDataParallel(model, device_ids=[local_rank] if torch.cuda.is_available() else None).to(device)
 
     # 3. SETUP LOSS AND OPTIMIZER
     criterion = torch.nn.MSELoss()
@@ -55,7 +55,7 @@ def run(self, main_address: str, main_port: int, num_nodes: int, node_rank: int)
         )
 
 
-# 32 GPUs: (8 nodes x 4 v 100)
+# 32 GPUs: (2 nodes x 4 v 100)
 compute = L.CloudCompute("gpu-fast-multi")  # 4xV100
 component = MultiNode(PyTorchDistributed, num_nodes=2, cloud_compute=compute)
 app = L.LightningApp(component)

src/lightning_app/components/training.py

@@ -147,33 +147,28 @@ def __init__(
                 the ServableModule API
         """
         super().__init__()
-        self.ws = structures.List()
-        self.has_initialized = False
         self.script_path = script_path
         self.script_args = script_args
         self.num_nodes = num_nodes
-        self._cloud_compute = cloud_compute  # TODO: Add support for cloudCompute
         self.sanity_serving = sanity_serving
         self._script_runner = script_runner
         self._script_runner_kwargs = script_runner_kwargs
 
-    def run(self, **run_kwargs):
-        if not self.has_initialized:
-            for node_rank in range(self.num_nodes):
-                self.ws.append(
-                    self._script_runner(
-                        script_path=self.script_path,
-                        script_args=self.script_args,
-                        cloud_compute=self._cloud_compute,
-                        node_rank=node_rank,
-                        sanity_serving=self.sanity_serving,
-                        num_nodes=self.num_nodes,
-                        **self._script_runner_kwargs,
-                    )
+        self.ws = structures.List()
+        for node_rank in range(self.num_nodes):
+            self.ws.append(
+                self._script_runner(
+                    script_path=self.script_path,
+                    script_args=self.script_args,
+                    cloud_compute=cloud_compute,
+                    node_rank=node_rank,
+                    sanity_serving=self.sanity_serving,
+                    num_nodes=self.num_nodes,
+                    **self._script_runner_kwargs,
                 )
+            )
 
-            self.has_initialized = True
-
+    def run(self, **run_kwargs):
         for work in self.ws:
             if all(w.internal_ip for w in self.ws):
                 internal_urls = [(w.internal_ip, w.port) for w in self.ws]