Parallel sky down (skypilot-org#659)

* fix multi-thread

* refactor

* Address comment

* format

* hidden variable

* Progress bar for termination

* fix

* format

* mitigate logging problem

* rename

Author: Michaelvll

sky/backends/backend_utils.py

@@ -11,6 +11,7 @@
 import subprocess
 import sys
 import textwrap
+import threading
 import time
 import typing
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -1188,3 +1189,20 @@ def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         code = ';'.join(code)
         return f'python3 -u -c {code!r}'
+
+
+class NoOpConsole:
+    """An empty class for multi-threaded console.status."""
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+def safe_console_status(msg: str):
+    """A wrapper for multi-threaded console.status."""
+    if threading.current_thread() is threading.main_thread():
+        return console.status(msg)
+    return NoOpConsole()

sky/backends/cloud_vm_ray_backend.py

@@ -20,7 +20,6 @@
 
 import colorama
 import filelock
-from rich import console as rich_console
 
 import sky
 from sky import backends
@@ -48,7 +47,6 @@
 SKY_REMOTE_RAY_VERSION = backend_utils.SKY_REMOTE_RAY_VERSION
 
 logger = sky_logging.init_logger(__name__)
-console = rich_console.Console()
 
 _PATH_SIZE_MEGABYTES_WARN_THRESHOLD = 256
 
@@ -726,8 +724,9 @@ def _try_provision_tpu(self, to_provision: 'resources_lib.Resources',
         assert 'tpu-create-script' in config_dict, \
             'Expect TPU provisioning with gcloud.'
         try:
-            with console.status('[bold cyan]Provisioning TPU '
-                                f'[green]{tpu_name}[/]'):
+            with backend_utils.safe_console_status(
+                    '[bold cyan]Provisioning TPU '
+                    f'[green]{tpu_name}[/]'):
                 backend_utils.run(f'bash {config_dict["tpu-create-script"]}',
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
@@ -1317,7 +1316,8 @@ def provision(self,
                 # PENDING / RUNNING jobs for the real status, since we do not
                 # know the actual previous status of the cluster.
                 cmd = backend_utils.JobLibCodeGen.update_status()
-                with console.status('[bold cyan]Preparing Job Queue'):
+                with backend_utils.safe_console_status(
+                        '[bold cyan]Preparing Job Queue'):
                     returncode, _, stderr = self.run_on_head(
                         handle, cmd, require_outputs=True)
                 backend_utils.handle_returncode(returncode, cmd,
@@ -1403,7 +1403,7 @@ def _sync_workdir_node(ip):
         tail_cmd = f'tail -n100 -f {log_path}'
         logger.info('To view detailed progress: '
                     f'{style.BRIGHT}{tail_cmd}{style.RESET_ALL}')
-        with console.status('[bold cyan]Syncing[/]'):
+        with backend_utils.safe_console_status('[bold cyan]Syncing[/]'):
             backend_utils.run_in_parallel(_sync_workdir_node, ip_list)
 
     def sync_file_mounts(
@@ -1478,7 +1478,7 @@ def _sync_node(ip):
             logger.info(f'{fore.CYAN}Syncing (to {num_nodes} node{plural}): '
                         f'{style.BRIGHT}{src}{style.RESET_ALL} -> '
                         f'{style.BRIGHT}{dst}{style.RESET_ALL}')
-            with console.status('[bold cyan]Syncing[/]'):
+            with backend_utils.safe_console_status('[bold cyan]Syncing[/]'):
                 backend_utils.run_in_parallel(_sync_node, ip_list)
 
         # Check the files and warn
@@ -1633,7 +1633,8 @@ def _setup_node(ip: int) -> int:
             plural = 's' if num_nodes > 1 else ''
             logger.info(f'{fore.CYAN}Running setup on {num_nodes} node{plural}.'
                         f'{style.RESET_ALL}')
-            with console.status('[bold cyan]Running setup[/]'):
+            with backend_utils.safe_console_status(
+                    '[bold cyan]Running setup[/]'):
                 backend_utils.run_in_parallel(_setup_node, ip_list)
         logger.info(f'{fore.GREEN}Setup completed.{style.RESET_ALL}')
         end = time.time()
@@ -1989,8 +1990,8 @@ def teardown_no_lock(self,
             # autoscaler.
             resource_group = config['provider']['resource_group']
             terminate_cmd = f'az group delete -y --name {resource_group}'
-            with console.status(f'[bold cyan]Terminating '
-                                f'[green]{cluster_name}'):
+            with backend_utils.safe_console_status(f'[bold cyan]Terminating '
+                                                   f'[green]{cluster_name}'):
                 returncode, stdout, stderr = log_lib.run_with_log(
                     terminate_cmd,
                     log_abs_path,
@@ -2014,8 +2015,9 @@ def teardown_no_lock(self,
                 terminate_cmd = (
                     f'aws ec2 terminate-instances --region {region} '
                     f'--instance-ids $({query_cmd})')
-                with console.status(f'[bold cyan]Terminating '
-                                    f'[green]{cluster_name}'):
+                with backend_utils.safe_console_status(
+                        f'[bold cyan]Terminating '
+                        f'[green]{cluster_name}'):
                     returncode, stdout, stderr = log_lib.run_with_log(
                         terminate_cmd,
                         log_abs_path,
@@ -2031,8 +2033,9 @@ def teardown_no_lock(self,
                 terminate_cmd = (
                     f'gcloud compute instances delete --zone={zone} --quiet '
                     f'$({query_cmd})')
-                with console.status(f'[bold cyan]Terminating '
-                                    f'[green]{cluster_name}'):
+                with backend_utils.safe_console_status(
+                        f'[bold cyan]Terminating '
+                        f'[green]{cluster_name}'):
                     returncode, stdout, stderr = log_lib.run_with_log(
                         terminate_cmd,
                         log_abs_path,
@@ -2052,16 +2055,18 @@ def teardown_no_lock(self,
                 f.flush()
 
                 teardown_verb = 'Terminating' if terminate else 'Stopping'
-                with console.status(f'[bold cyan]{teardown_verb} '
-                                    f'[green]{cluster_name}'):
+                with backend_utils.safe_console_status(
+                        f'[bold cyan]{teardown_verb} '
+                        f'[green]{cluster_name}'):
                     returncode, stdout, stderr = log_lib.run_with_log(
                         ['ray', 'down', '-y', f.name],
                         log_abs_path,
                         stream_logs=False,
                         require_outputs=True)
 
             if handle.tpu_delete_script is not None:
-                with console.status('[bold cyan]Terminating TPU...'):
+                with backend_utils.safe_console_status(
+                        '[bold cyan]Terminating TPU...'):
                     tpu_rc, tpu_stdout, tpu_stderr = log_lib.run_with_log(
                         ['bash', handle.tpu_delete_script],
                         log_abs_path,
@@ -2098,7 +2103,7 @@ def teardown_no_lock(self,
         backend_utils.SSHConfigHelper.remove_cluster(cluster_name,
                                                      handle.head_ip,
                                                      auth_config)
-        name = global_user_state.get_cluster_name_from_handle(handle)
+        name = handle.cluster_name
         global_user_state.remove_cluster(name, terminate=terminate)
 
         if terminate:

sky/cli.py

@@ -40,7 +40,7 @@
 
 import click
 import pendulum
-from rich import console as rich_console
+from rich import progress as rich_progress
 
 import sky
 from sky import backends
@@ -59,7 +59,6 @@
     from sky.backends import backend as backend_lib
 
 logger = sky_logging.init_logger(__name__)
-console = rich_console.Console()
 
 _CLUSTER_FLAG_HELP = """\
 A cluster name. If provided, either reuse an existing cluster with that name or
@@ -1302,32 +1301,46 @@ def _terminate_or_stop_clusters(names: Tuple[str],
             abort=True,
             show_default=True)
 
-    for record in to_down:  # TODO: parallelize.
+    progress = rich_progress.Progress(transient=True)
+    operation = 'Terminating' if terminate else 'Stopping'
+    plural = 's' if len(to_down) > 1 else ''
+    task = progress.add_task(
+        f'[bold cyan]{operation} {len(to_down)} cluster{plural}[/]',
+        total=len(to_down))
+    progress.start()
+
+    def _terminate_or_stop(record):
         name = record['name']
         handle = record['handle']
         backend = backend_utils.get_backend_from_handle(handle)
         if (isinstance(backend, backends.CloudVmRayBackend) and
                 handle.launched_resources.use_spot and not terminate):
+            # Disable spot instances to be stopped.
             # TODO(suquark): enable GCP+spot to be stopped in the future.
+            progress.stop()
             click.secho(
                 f'Stopping cluster {name}... skipped, because spot instances '
                 'may lose attached volumes. ',
                 fg='green')
             click.echo('  To terminate the cluster, run: ', nl=False)
             click.secho(f'sky down {name}', bold=True)
-            continue
-        success = backend.teardown(handle, terminate=terminate, purge=purge)
-        operation = 'Terminating' if terminate else 'Stopping'
-        if success:
-            click.secho(f'{operation} cluster {name}...done.', fg='green')
-            if not terminate:
-                click.echo('  To restart the cluster, run: ', nl=False)
-                click.secho(f'sky start {name}', bold=True)
         else:
-            click.secho(
-                f'{operation} cluster {name}...failed. '
-                'Please check the logs and try again.',
-                fg='red')
+            success = backend.teardown(handle, terminate=terminate, purge=purge)
+            progress.stop()
+            if success:
+                click.secho(f'{operation} cluster {name}...done.', fg='green')
+                if not terminate:
+                    click.echo('  To restart the cluster, run: ', nl=False)
+                    click.secho(f'sky start {name}', bold=True)
+            else:
+                click.secho(
+                    f'{operation} cluster {name}...failed. '
+                    'Please check the logs and try again.',
+                    fg='red')
+        progress.update(task, advance=1)
+        progress.start()
+
+    backend_utils.run_in_parallel(_terminate_or_stop, to_down)
 
 
 @_interactive_node_cli_command