Improve UX for autostopping (skypilot-org#676)

* Add progress bar for status refreshing

* Keep autostop after refreshing

* Add glob for start

* Fix message for autostop

* Fix messages for autostop

Author: Michaelvll

sky/backends/backend_utils.py

@@ -20,6 +20,7 @@
 
 import jinja2
 import rich.console as rich_console
+import rich.progress as rich_progress
 
 import sky
 from sky import authentication as auth
@@ -1140,7 +1141,12 @@ def get_clusters(refresh: bool) -> List[Dict[str, Any]]:
     records = global_user_state.get_clusters()
     if not refresh:
         return records
-    return [_ping_cluster_or_set_to_stopped(record) for record in records]
+    updated_records = []
+    for record in rich_progress.track(records,
+                                      description='Refreshing cluster status'):
+        record = _ping_cluster_or_set_to_stopped(record)
+        updated_records.append(record)
+    return updated_records
 
 
 def query_head_ip_with_retries(cluster_yaml: str, retry_count: int = 1) -> str:

sky/cli.py

@@ -1114,7 +1114,7 @@ def stop(
               default=None,
               is_flag=True,
               help='Tear down all existing clusters.')
-@click.option('--idle_minutes',
+@click.option('--idle-minutes',
               '-i',
               type=int,
               default=None,
@@ -1215,6 +1215,13 @@ def _filter(name, all_clusters):
                     return cluster_record
             return None
 
+        # Get GLOB cluster names
+        glob_clusters = []
+        for cluster in clusters:
+            glob_cluster = global_user_state.get_glob_cluster_names(cluster)
+            glob_clusters.extend(glob_cluster)
+        clusters = list(set(glob_clusters))
+
         all_clusters = global_user_state.get_clusters()
         for name in clusters:
             record = _filter(name, all_clusters)
@@ -1407,7 +1414,8 @@ def _terminate_or_stop_clusters(
 
     operation = 'Terminating' if terminate else 'Stopping'
     if idle_minutes_to_autostop is not None:
-        operation = 'Scheduling auto-stop on'
+        verb = 'Scheduling' if idle_minutes_to_autostop >= 0 else 'Cancelling'
+        operation = f'{verb} auto-stop on'
     plural = 's' if len(to_down) > 1 else ''
     progress = rich_progress.Progress(transient=True)
     task = progress.add_task(
@@ -1431,34 +1439,32 @@ def _terminate_or_stop(record):
                 f'{colorama.Style.RESET_ALL}')
         elif idle_minutes_to_autostop is not None:
             cluster_status = backend_utils.get_status_from_cluster_name(name)
-            if cluster_status != global_user_state.ClusterStatus.UP:
-                message = (
-                    f'{colorama.Fore.GREEN}Scheduling autostop for cluster '
-                    f'{name} (status: {cluster_status.value})... skipped'
-                    f'{colorama.Style.RESET_ALL}\n'
-                    '  Auto-stop can only be scheduled on '
-                    f'{global_user_state.ClusterStatus.UP.value} cluster.')
-            elif not isinstance(backend, backends.CloudVmRayBackend):
-                message = (
-                    f'{colorama.Fore.GREEN}Scheduling auto-stop for cluster '
-                    f'{name}... skipped{colorama.Style.RESET_ALL}\n'
-                    '  Auto-stopping is only supported by backend: '
-                    f'{backends.CloudVmRayBackend.NAME}')
+            if not isinstance(backend, backends.CloudVmRayBackend):
+                message = (f'{colorama.Fore.GREEN}{operation} cluster '
+                           f'{name}... skipped{colorama.Style.RESET_ALL}'
+                           '\n  Auto-stopping is only supported by backend: '
+                           f'{backends.CloudVmRayBackend.NAME}')
             else:
-                backend.set_autostop(handle, idle_minutes_to_autostop)
-                if idle_minutes_to_autostop < 0:
+                if cluster_status != global_user_state.ClusterStatus.UP:
                     message = (
-                        f'{colorama.Fore.GREEN}Cancelling auto-stop for '
-                        f'cluster {name}...done{colorama.Style.RESET_ALL}')
+                        f'{colorama.Fore.GREEN}{operation} cluster '
+                        f'{name} (status: {cluster_status.value})... skipped'
+                        f'{colorama.Style.RESET_ALL}'
+                        '\n  Auto-stop can only be run on '
+                        f'{global_user_state.ClusterStatus.UP.value} cluster.')
                 else:
+                    backend.set_autostop(handle, idle_minutes_to_autostop)
                     message = (
-                        f'{colorama.Fore.GREEN}Scheduling auto-stop for '
-                        f'cluster {name}...done{colorama.Style.RESET_ALL}\n'
-                        f'  The cluster will be stopped after '
-                        f'{idle_minutes_to_autostop} minutes of idleness.\n'
-                        '  To cancel the autostop, run: '
-                        f'{colorama.Style.BRIGHT}sky autostop {name} --cancel'
-                        f'{colorama.Style.RESET_ALL}')
+                        f'{colorama.Fore.GREEN}{operation} '
+                        f'cluster {name}...done{colorama.Style.RESET_ALL}')
+                    if idle_minutes_to_autostop >= 0:
+                        message += (
+                            f'\n  The cluster will be stopped after '
+                            f'{idle_minutes_to_autostop} minutes of idleness.'
+                            '\n  To cancel the autostop, run: '
+                            f'{colorama.Style.BRIGHT}'
+                            f'sky autostop {name} --cancel'
+                            f'{colorama.Style.RESET_ALL}')
         else:
             success = backend.teardown(handle, terminate=terminate, purge=purge)
             if success:

sky/global_user_state.py

@@ -164,9 +164,11 @@ def add_or_update_cluster(cluster_name: str,
         'VALUES (?, ?, ?, ?, ?, '
         # Keep the old autostop value if it exists, otherwise set it to
         # default -1.
-        'COALESCE((SELECT autostop FROM clusters WHERE name=?), -1))',
+        'COALESCE('
+        '(SELECT autostop FROM clusters WHERE name=? AND status!=?), -1)'
+        ')',  # VALUES
         (cluster_name, cluster_launched_at, handle, last_use, status.value,
-         cluster_name))
+         cluster_name, ClusterStatus.STOPPED.value))
     _DB.conn.commit()
 
 
@@ -190,11 +192,10 @@ def remove_cluster(cluster_name: str, terminate: bool):
         # will directly try to ssh, which leads to timeout.
         handle.head_ip = None
         _DB.cursor.execute(
-            'UPDATE clusters SET handle=(?), status=(?), autostop=(?) '
+            'UPDATE clusters SET handle=(?), status=(?) '
             'WHERE name=(?)', (
                 pickle.dumps(handle),
                 ClusterStatus.STOPPED.value,
-                -1,
                 cluster_name,
             ))
     _DB.conn.commit()