Repository: skypilot-org/skypilot
Pull request: Parallel sky down
Closes #641. This PR makes `sky down` run in parallel. It also mitigates the problem with #656 to tear down the 100 clusters faster.

Main changes:
1. Make the database and rich console compatible with multi-thread.

Tested:
- [x] `for i in {1..4}; do sky launch -c min-$i ./examples/minimal.yaml -y 2>&1 | tee $i-launch.log &; done; wait; sky down -a -y`
- [x] `sky down` clusters that contain some failure. 

Problem statement:
Can sky down be parallelized?
I created 20 clusters on Azure for testing purpose.
After testing, I'd like to `sky down -a` but terminating each Azure cluster is taking >5mins.
In total 100mins is needed to clean up all clusters.
Is it possible to send all the termination signals in parallel?

Hints from issue discussion:
Note: I just tried it with `backend_utils.run_in_parallel`, it seems that the `sqlite3` cannot be run in multiple threads. We may need to try multiple processes or ray actors.
we could also just use `conn = sqlite3.connect('your.db', check_same_thread=False)` to allow sqlite3 on multiple threads then create our own locks to handle synchronization
With the launch pattern in #656, parallel teardown will be very helpful.

Commit message: format

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 76ed4441b6c..615f80adc37 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -11,6 +11,7 @@
 import subprocess
 import sys
 import textwrap
+import threading
 import time
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import uuid
@@ -1190,3 +1191,24 @@ def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         code = ';'.join(code)
         return f'python3 -u -c {code!r}'
+
+
+class ThreadConsole:
+    """An empty class for multi-threaded console.status."""
+    PATTERN = re.compile(r'\[.*?\]')
+
+    def __init__(self, msg: str):
+        self.msg = self.PATTERN.sub('', msg)
+
+    def __enter__(self):
+        print(self.msg)
+
+    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 ThreadConsole(msg)
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 86be1517933..813e26b178d 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -18,7 +18,6 @@
 
 import colorama
 import filelock
-from rich import console as rich_console
 
 import sky
 from sky import backends
@@ -47,7 +46,6 @@
 SKYLET_REMOTE_PATH = backend_utils.SKY_REMOTE_PATH
 
 logger = sky_logging.init_logger(__name__)
-console = rich_console.Console()
 
 _PATH_SIZE_MEGABYTES_WARN_THRESHOLD = 256
 
@@ -721,8 +719,9 @@ def _try_provision_tpu(self, to_provision: 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)
@@ -1298,7 +1297,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,
@@ -1384,7 +1384,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(
@@ -1459,7 +1459,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
@@ -1614,7 +1614,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()
@@ -1972,8 +1973,8 @@ def _teardown(self, handle: ResourceHandle, terminate: bool) -> bool:
             # 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,
@@ -1997,8 +1998,9 @@ def _teardown(self, handle: ResourceHandle, terminate: bool) -> bool:
                 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,
@@ -2014,8 +2016,9 @@ def _teardown(self, handle: ResourceHandle, terminate: bool) -> bool:
                 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,
@@ -2035,8 +2038,9 @@ def _teardown(self, handle: ResourceHandle, terminate: bool) -> bool:
                 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,
@@ -2044,7 +2048,8 @@ def _teardown(self, handle: ResourceHandle, terminate: bool) -> bool:
                         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,
@@ -2071,7 +2076,7 @@ def _teardown(self, handle: ResourceHandle, terminate: bool) -> bool:
         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:
diff --git a/sky/cli.py b/sky/cli.py
index 8d57751d0ac..489a17666bc 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -1287,7 +1287,7 @@ def _terminate_or_stop_clusters(names: Tuple[str], apply_to_all: Optional[bool],
             abort=True,
             show_default=True)
 
-    for record in to_down:  # TODO: parallelize.
+    def _terminate_or_stop(record):
         name = record['name']
         handle = record['handle']
         backend = backend_utils.get_backend_from_handle(handle)
@@ -1300,7 +1300,7 @@ def _terminate_or_stop_clusters(names: Tuple[str], apply_to_all: Optional[bool],
                 fg='green')
             click.echo('  To terminate the cluster, run: ', nl=False)
             click.secho(f'sky down {name}', bold=True)
-            continue
+            return
         success = backend.teardown(handle, terminate=terminate)
         operation = 'Terminating' if terminate else 'Stopping'
         if success:
@@ -1314,6 +1314,8 @@ def _terminate_or_stop_clusters(names: Tuple[str], apply_to_all: Optional[bool],
                 'Please check the logs and try again.',
                 fg='red')
 
+    backend_utils.run_in_parallel(_terminate_or_stop, to_down)
+
 
 @_interactive_node_cli_command
 def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 6ca2fb23c27..38e11a483e2 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -13,6 +13,7 @@
 import pickle
 import sqlite3
 import sys
+import threading
 import time
 from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
@@ -27,31 +28,42 @@
 _DB_PATH = os.path.expanduser('~/.sky/state.db')
 os.makedirs(pathlib.Path(_DB_PATH).parents[0], exist_ok=True)
 
-_CONN = sqlite3.connect(_DB_PATH)
-_CURSOR = _CONN.cursor()
-
-# Table for Clusters
-_CURSOR.execute("""\
-    CREATE TABLE IF NOT EXISTS clusters (
-    name TEXT PRIMARY KEY,
-    lauched_at INTEGER,
-    handle BLOB,
-    last_use TEXT,
-    status TEXT)""")
-# Table for Sky Config (e.g. enabled clouds)
-_CURSOR.execute("""\
-    CREATE TABLE IF NOT EXISTS config (
-    key TEXT PRIMARY KEY, value TEXT)""")
-# Table for Storage
-_CURSOR.execute("""\
-    CREATE TABLE IF NOT EXISTS storage (
-    name TEXT PRIMARY KEY,
-    lauched_at INTEGER,
-    handle BLOB,
-    last_use TEXT,
-    status TEXT)""")
-
-_CONN.commit()
+
+class SQLiteConn(threading.local):
+    """Thread-local connection to the sqlite3 database."""
+
+    def __init__(self):
+        super().__init__()
+        self.conn = sqlite3.connect(_DB_PATH)
+        self.cursor = self.conn.cursor()
+        self.create_table()
+
+    def create_table(self):
+        # Table for Clusters
+        self.cursor.execute("""\
+            CREATE TABLE IF NOT EXISTS clusters (
+            name TEXT PRIMARY KEY,
+            lauched_at INTEGER,
+            handle BLOB,
+            last_use TEXT,
+            status TEXT)""")
+        # Table for Sky Config (e.g. enabled clouds)
+        self.cursor.execute("""\
+            CREATE TABLE IF NOT EXISTS config (
+            key TEXT PRIMARY KEY, value TEXT)""")
+        # Table for Storage
+        self.cursor.execute("""\
+            CREATE TABLE IF NOT EXISTS storage (
+            name TEXT PRIMARY KEY,
+            lauched_at INTEGER,
+            handle BLOB,
+            last_use TEXT,
+            status TEXT)""")
+
+        self.conn.commit()
+
+
+_DB = SQLiteConn()
 
 
 class ClusterStatus(enum.Enum):
@@ -122,23 +134,24 @@ def add_or_update_cluster(cluster_name: str,
     handle = pickle.dumps(cluster_handle)
     last_use = _get_pretty_entry_point()
     status = ClusterStatus.UP if ready else ClusterStatus.INIT
-    _CURSOR.execute(
+    _DB.cursor.execute(
         'INSERT OR REPLACE INTO clusters VALUES (?, ?, ?, ?, ?)',
         (cluster_name, cluster_launched_at, handle, last_use, status.value))
-    _CONN.commit()
+    _DB.conn.commit()
 
 
 def update_last_use(cluster_name: str):
     """Updates the last used command for the cluster."""
-    _CURSOR.execute('UPDATE clusters SET last_use=(?) WHERE name=(?)',
-                    (_get_pretty_entry_point(), cluster_name))
-    _CONN.commit()
+    _DB.cursor.execute('UPDATE clusters SET last_use=(?) WHERE name=(?)',
+                       (_get_pretty_entry_point(), cluster_name))
+    _DB.conn.commit()
 
 
 def remove_cluster(cluster_name: str, terminate: bool):
     """Removes cluster_name mapping."""
     if terminate:
-        _CURSOR.execute('DELETE FROM clusters WHERE name=(?)', (cluster_name,))
+        _DB.cursor.execute('DELETE FROM clusters WHERE name=(?)',
+                           (cluster_name,))
     else:
         handle = get_handle_from_cluster_name(cluster_name)
         if handle is None:
@@ -146,62 +159,62 @@ def remove_cluster(cluster_name: str, terminate: bool):
         # Must invalidate head_ip: otherwise 'sky cpunode' on a stopped cpunode
         # will directly try to ssh, which leads to timeout.
         handle.head_ip = None
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE clusters SET handle=(?), status=(?) WHERE name=(?)', (
                 pickle.dumps(handle),
                 ClusterStatus.STOPPED.value,
                 cluster_name,
             ))
-    _CONN.commit()
+    _DB.conn.commit()
 
 
 def get_handle_from_cluster_name(
         cluster_name: str) -> Optional['backends.Backend.ResourceHandle']:
     assert cluster_name is not None, 'cluster_name cannot be None'
-    rows = _CURSOR.execute('SELECT handle FROM clusters WHERE name=(?)',
-                           (cluster_name,))
+    rows = _DB.cursor.execute('SELECT handle FROM clusters WHERE name=(?)',
+                              (cluster_name,))
     for (handle,) in rows:
         return pickle.loads(handle)
 
 
 def get_glob_cluster_names(cluster_name: str) -> List[str]:
     assert cluster_name is not None, 'cluster_name cannot be None'
-    rows = _CURSOR.execute('SELECT name FROM clusters WHERE name GLOB (?)',
-                           (cluster_name,))
+    rows = _DB.cursor.execute('SELECT name FROM clusters WHERE name GLOB (?)',
+                              (cluster_name,))
     return [row[0] for row in rows]
 
 
 def get_cluster_name_from_handle(
         cluster_handle: 'backends.Backend.ResourceHandle') -> Optional[str]:
     handle = pickle.dumps(cluster_handle)
-    rows = _CURSOR.execute('SELECT name FROM clusters WHERE handle=(?)',
-                           (handle,))
+    rows = _DB.cursor.execute('SELECT name FROM clusters WHERE handle=(?)',
+                              (handle,))
     for (name,) in rows:
         return name
 
 
 def get_status_from_cluster_name(
         cluster_name: Optional[str]) -> Optional[ClusterStatus]:
-    rows = _CURSOR.execute('SELECT status FROM clusters WHERE name=(?)',
-                           (cluster_name,))
+    rows = _DB.cursor.execute('SELECT status FROM clusters WHERE name=(?)',
+                              (cluster_name,))
     for (status,) in rows:
         return ClusterStatus[status]
 
 
 def set_cluster_status(cluster_name: str, status: ClusterStatus) -> None:
-    _CURSOR.execute('UPDATE clusters SET status=(?) WHERE name=(?)', (
+    _DB.cursor.execute('UPDATE clusters SET status=(?) WHERE name=(?)', (
         status.value,
         cluster_name,
     ))
-    count = _CURSOR.rowcount
-    _CONN.commit()
+    count = _DB.cursor.rowcount
+    _DB.conn.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
 def get_clusters() -> List[Dict[str, Any]]:
-    rows = _CURSOR.execute('select * from clusters')
+    rows = _DB.cursor.execute('select * from clusters')
     records = []
     for name, launched_at, handle, last_use, status in rows:
         # TODO: use namedtuple instead of dict
@@ -216,8 +229,8 @@ def get_clusters() -> List[Dict[str, Any]]:
 
 
 def get_enabled_clouds() -> List[clouds.Cloud]:
-    rows = _CURSOR.execute('SELECT value FROM config WHERE key = ?',
-                           (_ENABLED_CLOUDS_KEY,))
+    rows = _DB.cursor.execute('SELECT value FROM config WHERE key = ?',
+                              (_ENABLED_CLOUDS_KEY,))
     ret = []
     for (value,) in rows:
         ret = json.loads(value)
@@ -226,9 +239,9 @@ def get_enabled_clouds() -> List[clouds.Cloud]:
 
 
 def set_enabled_clouds(enabled_clouds: List[str]) -> None:
-    _CURSOR.execute('INSERT OR REPLACE INTO config VALUES (?, ?)',
-                    (_ENABLED_CLOUDS_KEY, json.dumps(enabled_clouds)))
-    _CONN.commit()
+    _DB.cursor.execute('INSERT OR REPLACE INTO config VALUES (?, ?)',
+                       (_ENABLED_CLOUDS_KEY, json.dumps(enabled_clouds)))
+    _DB.conn.commit()
 
 
 def add_or_update_storage(storage_name: str,
@@ -244,25 +257,25 @@ def status_check(status):
     if not status_check(storage_status):
         raise ValueError(f'Error in updating global state. Storage Status '
                          f'{storage_status} is passed in incorrectly')
-    _CURSOR.execute('INSERT OR REPLACE INTO storage VALUES (?, ?, ?, ?, ?)',
-                    (storage_name, storage_launched_at, handle, last_use,
-                     storage_status.value))
-    _CONN.commit()
+    _DB.cursor.execute('INSERT OR REPLACE INTO storage VALUES (?, ?, ?, ?, ?)',
+                       (storage_name, storage_launched_at, handle, last_use,
+                        storage_status.value))
+    _DB.conn.commit()
 
 
 def remove_storage(storage_name: str):
     """Removes Storage from Database"""
-    _CURSOR.execute('DELETE FROM storage WHERE name=(?)', (storage_name,))
-    _CONN.commit()
+    _DB.cursor.execute('DELETE FROM storage WHERE name=(?)', (storage_name,))
+    _DB.conn.commit()
 
 
 def set_storage_status(storage_name: str, status: StorageStatus) -> None:
-    _CURSOR.execute('UPDATE storage SET status=(?) WHERE name=(?)', (
+    _DB.cursor.execute('UPDATE storage SET status=(?) WHERE name=(?)', (
         status.value,
         storage_name,
     ))
-    count = _CURSOR.rowcount
-    _CONN.commit()
+    count = _DB.cursor.rowcount
+    _DB.conn.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Storage{storage_name} not found.')
@@ -270,19 +283,19 @@ def set_storage_status(storage_name: str, status: StorageStatus) -> None:
 
 def get_storage_status(storage_name: str) -> None:
     assert storage_name is not None, 'storage_name cannot be None'
-    rows = _CURSOR.execute('SELECT status FROM storage WHERE name=(?)',
-                           (storage_name,))
+    rows = _DB.cursor.execute('SELECT status FROM storage WHERE name=(?)',
+                              (storage_name,))
     for (status,) in rows:
         return StorageStatus[status]
 
 
 def set_storage_handle(storage_name: str, handle: 'Storage.StorageMetadata'):
-    _CURSOR.execute('UPDATE storage SET handle=(?) WHERE name=(?)', (
+    _DB.cursor.execute('UPDATE storage SET handle=(?) WHERE name=(?)', (
         pickle.dumps(handle),
         storage_name,
     ))
-    count = _CURSOR.rowcount
-    _CONN.commit()
+    count = _DB.cursor.rowcount
+    _DB.conn.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Storage{storage_name} not found.')
@@ -290,14 +303,14 @@ def set_storage_handle(storage_name: str, handle: 'Storage.StorageMetadata'):
 
 def get_handle_from_storage_name(storage_name: str):
     assert storage_name is not None, 'storage_name cannot be None'
-    rows = _CURSOR.execute('SELECT handle FROM storage WHERE name=(?)',
-                           (storage_name,))
+    rows = _DB.cursor.execute('SELECT handle FROM storage WHERE name=(?)',
+                              (storage_name,))
     for (handle,) in rows:
         return pickle.loads(handle)
 
 
 def get_storage() -> List[Dict[str, Any]]:
-    rows = _CURSOR.execute('select * from storage')
+    rows = _DB.cursor.execute('select * from storage')
     records = []
     for name, launched_at, handle, last_use, status in rows:
         # TODO: use namedtuple instead of dict


Code review comments:
diff_hunk: @@ -27,31 +28,42 @@
 _DB_PATH = os.path.expanduser('~/.sky/state.db')
 os.makedirs(pathlib.Path(_DB_PATH).parents[0], exist_ok=True)
 
-_CONN = sqlite3.connect(_DB_PATH)
-_CURSOR = _CONN.cursor()
-
-# Table for Clusters
-_CURSOR.execute("""\
-    CREATE TABLE IF NOT EXISTS clusters (
-    name TEXT PRIMARY KEY,
-    lauched_at INTEGER,
-    handle BLOB,
-    last_use TEXT,
-    status TEXT)""")
-# Table for Sky Config (e.g. enabled clouds)
-_CURSOR.execute("""\
-    CREATE TABLE IF NOT EXISTS config (
-    key TEXT PRIMARY KEY, value TEXT)""")
-# Table for Storage
-_CURSOR.execute("""\
-    CREATE TABLE IF NOT EXISTS storage (
-    name TEXT PRIMARY KEY,
-    lauched_at INTEGER,
-    handle BLOB,
-    last_use TEXT,
-    status TEXT)""")
-
-_CONN.commit()
+
+class SQLiteConn(threading.local):
path: sky/global_user_state.py
text: @user1:
Minor: _SQLiteConn + _create_table()?

@author:
Good point! Fixed. Thanks!