Visitar URL original
Add async create_pools_to / delete_pools_above for elastic scaling · apple/axlearn@bb4d031 · GitHub
Skip to content

Commit bb4d031

Browse files
Alexandre Jameswangkuiyi
authored andcommitted
Add async create_pools_to / delete_pools_above for elastic scaling
GitOrigin-RevId: 937f46d349f897cbb02c835bf6aeb8903db12002
1 parent c2da3b2 commit bb4d031

3 files changed

Lines changed: 368 additions & 35 deletions

File tree

‎axlearn/cloud/gcp/node_pool.py‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,7 @@ def delete_node_pools(
329329
cluster: str,
330330
retry_interval: int = 30,
331331
wait_timeout: int = 0,
332-
):
332+
) -> int:
333333
"""Delete node pools.
334334
335335
See also:
@@ -344,6 +344,11 @@ def delete_node_pools(
344344
wait_timeout: Seconds to wait for node pools to delete.
345345
This function doesn't wait if wait_timeout <= 0.
346346
347+
Returns:
348+
Number of node pools still existing in `names` at the end of the call
349+
(i.e. pools whose status is not NOT_EXIST). When `wait_timeout > 0`
350+
and the call returns normally, this is always 0.
351+
347352
Raises:
348353
NodePoolDeletionError: If node pools are not deleted
349354
within wait_timeout if wait_timeout > 0.
@@ -382,14 +387,14 @@ def delete_node_pools(
382387
logging.info(
383388
"All %s node pools deletion took %s seconds", num_node_pools_deleted, elapsed_time
384389
)
385-
break
390+
return 0
386391

387392
if wait_timeout <= 0:
388393
logging.info(
389394
"Skip waiting for node pool deletion since wait_timeout (%s) is non-positive",
390395
wait_timeout,
391396
)
392-
break
397+
return len(names) - num_node_pools_deleted
393398

394399
elif elapsed_time > wait_timeout:
395400
raise NodePoolDeletionError(
@@ -511,7 +516,7 @@ def create_node_pools(
511516
enable_confidential_storage: Optional[bool] = None,
512517
retry_interval: int = 30,
513518
wait_timeout: int = 0,
514-
):
519+
) -> int:
515520
"""Create node pools.
516521
517522
See also:
@@ -541,6 +546,11 @@ def create_node_pools(
541546
wait_timeout: Seconds to wait for node pools to delete.
542547
This function doesn't wait if wait_timeout <= 0.
543548
549+
Returns:
550+
Number of node pools in `names` whose status is RUNNING at the end of
551+
the call. When `wait_timeout > 0` and the call returns normally, this
552+
is always `len(names)`.
553+
544554
Raises:
545555
NodePoolCreationError: If node pools are not deleted
546556
within wait_timeout if wait_timeout > 0.
@@ -597,14 +607,14 @@ def create_node_pools(
597607
pre_provisioner_id,
598608
elapsed_time,
599609
)
600-
break
610+
return num_node_pools_running
601611

602612
if wait_timeout <= 0:
603613
logging.info(
604614
"Skip waiting for node pool creation since wait_timeout (%s) is non-positive",
605615
wait_timeout,
606616
)
607-
break
617+
return num_node_pools_running
608618

609619
elif elapsed_time > wait_timeout:
610620
raise NodePoolCreationError(

0 commit comments

Comments
 (0)