Skip to content

Commit abf0a40

Browse files
ethanwharrisBorda
authored andcommitted
[App] Improve LightningTrainerScript start-up time (#15751)
(cherry picked from commit c2c1974)
1 parent 9351d41 commit abf0a40

File tree

1 file changed

+13
-18
lines changed

1 file changed

+13
-18
lines changed

src/lightning_app/components/training.py

Lines changed: 13 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -147,33 +147,28 @@ def __init__(
147147
the ServableModule API
148148
"""
149149
super().__init__()
150-
self.ws = structures.List()
151-
self.has_initialized = False
152150
self.script_path = script_path
153151
self.script_args = script_args
154152
self.num_nodes = num_nodes
155-
self._cloud_compute = cloud_compute # TODO: Add support for cloudCompute
156153
self.sanity_serving = sanity_serving
157154
self._script_runner = script_runner
158155
self._script_runner_kwargs = script_runner_kwargs
159156

160-
def run(self, **run_kwargs):
161-
if not self.has_initialized:
162-
for node_rank in range(self.num_nodes):
163-
self.ws.append(
164-
self._script_runner(
165-
script_path=self.script_path,
166-
script_args=self.script_args,
167-
cloud_compute=self._cloud_compute,
168-
node_rank=node_rank,
169-
sanity_serving=self.sanity_serving,
170-
num_nodes=self.num_nodes,
171-
**self._script_runner_kwargs,
172-
)
157+
self.ws = structures.List()
158+
for node_rank in range(self.num_nodes):
159+
self.ws.append(
160+
self._script_runner(
161+
script_path=self.script_path,
162+
script_args=self.script_args,
163+
cloud_compute=cloud_compute,
164+
node_rank=node_rank,
165+
sanity_serving=self.sanity_serving,
166+
num_nodes=self.num_nodes,
167+
**self._script_runner_kwargs,
173168
)
169+
)
174170

175-
self.has_initialized = True
176-
171+
def run(self, **run_kwargs):
177172
for work in self.ws:
178173
if all(w.internal_ip for w in self.ws):
179174
internal_urls = [(w.internal_ip, w.port) for w in self.ws]

0 commit comments

Comments
 (0)