[Autostop/UX] Better UX when the cluster is partially stopped due to autostopping (skypilot-org#1637)

* Fix status logging warning

* Not reset autostop column when the cluster is autostopping

* Add a hint for INIT mode cluster which is autostopping

* set the indicator

* format

* print out payload

* fix 'to_down'

* Not error out for is_autostopping

* format

* address comments

* address comments

Author: Michaelvll

sky/backends/backend_utils.py

@@ -1950,24 +1950,34 @@ def _update_cluster_status_no_lock(
     #     abnormal.
     #
     # An abnormal cluster will transition to INIT and have any autostop setting
-    # reset.
+    # reset (unless it's autostopping/autodowning.).
     is_abnormal = ((0 < len(node_statuses) < handle.launched_nodes) or
                    any(status != global_user_state.ClusterStatus.STOPPED
                        for status in node_statuses))
     if is_abnormal:
-        # Reset the autostop to avoid false information with best effort.
-        # Side effect: if the status is refreshed during autostopping, the
-        # autostop field in the local cache will be reset, even though the
-        # cluster will still be correctly stopped.
-        try:
-            backend = backends.CloudVmRayBackend()
-            backend.set_autostop(handle, -1, stream_logs=False)
-        except (Exception, SystemExit) as e:  # pylint: disable=broad-except
-            logger.debug(f'Failed to reset autostop. Due to '
-                         f'{common_utils.format_exception(e)}')
-        global_user_state.set_cluster_autostop_value(handle.cluster_name,
-                                                     -1,
-                                                     to_down=False)
+        backend = get_backend_from_handle(handle)
+        if isinstance(backend,
+                      backends.CloudVmRayBackend) and record['autostop'] >= 0:
+            if not backend.is_definitely_autostopping(handle,
+                                                      stream_logs=False):
+                # Reset the autostopping as the cluster is abnormal, and may
+                # not correctly autostop. Resetting the autostop will let
+                # the user know that the autostop may not happen to avoid
+                # leakages from the assumption that the cluster will autostop.
+                try:
+                    backend.set_autostop(handle, -1, stream_logs=False)
+                except (Exception, SystemExit) as e:  # pylint: disable=broad-except
+                    logger.debug(f'Failed to reset autostop. Due to '
+                                 f'{common_utils.format_exception(e)}')
+                global_user_state.set_cluster_autostop_value(
+                    handle.cluster_name, -1, to_down=False)
+            else:
+                ux_utils.console_newline()
+                operation_str = 'autodowning' if record[
+                    'to_down'] else 'autostopping'
+                logger.info(
+                    f'Cluster {cluster_name!r} is {operation_str}. Setting to '
+                    'INIT status; try refresh again in a while.')
 
         # If the user starts part of a STOPPED cluster, we still need a status
         # to represent the abnormal status. For spot cluster, it can also

sky/backends/cloud_vm_ray_backend.py

@@ -3063,6 +3063,27 @@ def set_autostop(self,
             global_user_state.set_cluster_autostop_value(
                 handle.cluster_name, idle_minutes_to_autostop, down)
 
+    def is_definitely_autostopping(self,
+                                   handle: ResourceHandle,
+                                   stream_logs: bool = True) -> bool:
+        """Check if the cluster is autostopping.
+
+        Returns:
+            True if the cluster is definitely autostopping. It is possible
+            that the cluster is still autostopping when False is returned,
+            due to errors like transient network issues.
+        """
+        code = autostop_lib.AutostopCodeGen.is_autostopping()
+        returncode, stdout, stderr = self.run_on_head(handle,
+                                                      code,
+                                                      require_outputs=True,
+                                                      stream_logs=stream_logs)
+
+        if returncode == 0:
+            return common_utils.decode_payload(stdout)
+        logger.debug(f'Failed to check if cluster is autostopping: {stderr}')
+        return False
+
     # TODO(zhwu): Refactor this to a CommandRunner class, so different backends
     # can support its own command runner.
     @timeline.event

sky/skylet/autostop_lib.py

@@ -7,6 +7,7 @@
 
 from sky import sky_logging
 from sky.skylet import configs
+from sky.utils import common_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -16,6 +17,10 @@
 # user-issued commands (this module) and the Skylet process running the
 # AutostopEvent need to access that state.
 _AUTOSTOP_LAST_ACTIVE_TIME = 'autostop_last_active_time'
+# AutostopEvent sets this to the boot time when the autostop of the cluster
+# starts. This is used for checking whether the cluster is in the process
+# of autostopping for the current machine.
+_AUTOSTOP_INDICATOR = 'autostop_indicator'
 
 
 class AutostopConfig:
@@ -57,6 +62,26 @@ def set_autostop(idle_minutes: int, backend: Optional[str], down: bool) -> None:
         set_last_active_time_to_now()
 
 
+def set_autostopping_started() -> None:
+    """Sets the boot time of the machine when autostop starts.
+
+    This function should be called when the cluster is started to autostop,
+    and the boot time of the machine will be stored in the configs database
+    as an autostop indicator, which is used for checking whether the cluster
+    is in the process of autostopping. The indicator is valid only when the
+    machine has the same boot time as the one stored in the indicator.
+    """
+    logger.debug('Setting is_autostopping.')
+    configs.set_config(_AUTOSTOP_INDICATOR, str(psutil.boot_time()))
+
+
+def get_is_autostopping_payload() -> str:
+    """Returns whether the cluster is in the process of autostopping."""
+    result = configs.get_config(_AUTOSTOP_INDICATOR)
+    is_autostopping = (result == str(psutil.boot_time()))
+    return common_utils.encode_payload(is_autostopping)
+
+
 def get_last_active_time() -> float:
     """Returns the last active time, or -1 if none has been set."""
     result = configs.get_config(_AUTOSTOP_LAST_ACTIVE_TIME)
@@ -88,6 +113,11 @@ def set_autostop(cls, idle_minutes: int, backend: str, down: bool) -> str:
         ]
         return cls._build(code)
 
+    @classmethod
+    def is_autostopping(cls) -> str:
+        code = ['print(autostop_lib.get_is_autostopping_payload())']
+        return cls._build(code)
+
     @classmethod
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code

sky/skylet/configs.py

@@ -60,6 +60,7 @@ def get_config(key: str) -> Optional[str]:
         rows = cursor.execute('SELECT value FROM config WHERE key = ?', (key,))
         for (value,) in rows:
             return value
+        return None
 
 
 @ensure_table

sky/skylet/events.py

@@ -128,6 +128,7 @@ def _run(self):
     def _stop_cluster(self, autostop_config):
         if (autostop_config.backend ==
                 cloud_vm_ray_backend.CloudVmRayBackend.NAME):
+            autostop_lib.set_autostopping_started()
             self._replace_yaml_for_stopping(self._ray_yaml_path,
                                             autostop_config.down)