Repository: skypilot-org/skypilot
Pull request: [Autostop/UX] Better UX when the cluster is partially stopped due to autostopping
<!-- Describe the changes in this PR -->

Fixes #1633 
Fixes the confusion of cluster being in `INIT` status, when it is being autostopped and is partially stopped.

Changes:
1. keep the autostop column in the cluster table, when the cluster is partially stopped
2. Print a hint that the cluster is being autostopped when it is in INIT status

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Snippets from #1633
```
# launch multinode gcp cluster
sky launch -y -d -c test --num-nodes 2 --cloud gcp tests/test_yamls/minimal.yaml

# Try sky status --refresh. This works fine.
sky status test --refresh

# Schedule autostop
sky autostop -y test -i 0

sleep 60

# Try refreshing status - It should print out INIT status / hint for autostopping and autostop column.
sky status test --refresh
```
<img width="817" alt="image" src="https://user-images.githubusercontent.com/6753189/214964420-aa288ab3-2d51-4054-bb6a-502cc78f5f3c.png">


- [x] The following test
```
# launch multinode gcp cluster
sky launch -y -d -c test --num-nodes 2 --cloud gcp tests/test_yamls/minimal.yaml

# Try sky status --refresh. This works fine.
sky status test --refresh

# Schedule autostop
sky autostop -y test -i 5

# Manually stop the head node, no autostop output, no autostop column.
sky status -r
```
<img width="546" alt="image" src="https://user-images.githubusercontent.com/6753189/214966981-184e49e1-bc95-44d7-8ab2-b20065f4c54d.png">


Problem statement:
[Autostop] `sky status --refresh` after autostop fails for multi-node on GCP
## Reproduction
```
# launch multinode gcp cluster
sky launch -y -d -c test --num-nodes 2 --cloud gcp tests/test_yamls/minimal.yaml

# Try sky status --refresh. This works fine.
sky status test --refresh

# Schedule autostop
sky autostop -y test -i 1

# Wait for autostop
sleep 120

# Try refreshing status - fails
sky status test --refresh
```

## Log
```
gcpuser@ray-smoke-head-ff075d77-compute:/sky_repo$ sky status test --refresh
Refreshing status for 1 cluster ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━   0% -:--:--W 01-26 08:49:45 backend_utils.py:1346] Expected 1 worker IP(s); found 0: []
W 01-26 08:49:45 backend_utils.py:1346] This could happen if there is extra output from `ray get-worker-ips`, which should be inspected below.
W 01-26 08:49:45 backend_utils.py:1346] == Output ==
W 01-26 08:49:45 backend_utils.py:1346]
W 01-26 08:49:45 backend_utils.py:1346]
W 01-26 08:49:45 backend_utils.py:1346] == Output ends ==
Clusters
NAME  LAUNCHED    RESOURCES              STATUS   AUTOSTOP  COMMAND
test  5 mins ago  2x GCP(n2-standard-8)  STOPPED  1m        sky launch -y -d -c test ...

1 cluster has auto{stop,down} scheduled. Refresh statuses with: sky status --refresh
```

## Notes
* This happens for only a short duration (~1-2 min?) while the worker is stopping (checked in cloud console). Running `sky status test --refresh` afterwards works normally and shows status as stopped.
* This happens only when using autostop. I could not replicate by manually stopping a cluster `sky stop test` and then refresh status.
* This bug is also caught by `test_smoke.py::test_autostop`
* Not fully sure if this is related to multi-node or GCP - need to check if it also happens for single node and on different clouds.

Hints from issue discussion:
Just to add, in `test_autostop` the state is shown to be at `INIT` instead of `UP` or `STOPPED`, which causes the test to fail. [Logs here](https://gist.github.com/romilbhardwaj/4be928e440d914d66bccd1045d7a9e14).
Partially addressed by #1634 and #1638, removing P0 label

Commit message: address comments

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index b83dc5d8a99..6fc1b8b1a6b 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1950,24 +1950,35 @@ 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 is not going to 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
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 6ea6d4e7832..1f0e87a049e 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/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 network transiant 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
diff --git a/sky/skylet/autostop_lib.py b/sky/skylet/autostop_lib.py
index 0756ec29e14..39f824820b5 100644
--- a/sky/skylet/autostop_lib.py
+++ b/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,27 @@ def set_autostop(idle_minutes: int, backend: Optional[str], down: bool) -> None:
         set_last_active_time_to_now()
 
 
+def set_autostopping_indicator() -> 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.
+    It 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 +114,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
diff --git a/sky/skylet/configs.py b/sky/skylet/configs.py
index fadeed19428..b5d1776464f 100644
--- a/sky/skylet/configs.py
+++ b/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
diff --git a/sky/skylet/events.py b/sky/skylet/events.py
index d3b66dabfc0..e7e414c6199 100644
--- a/sky/skylet/events.py
+++ b/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_indicator()
             self._replace_yaml_for_stopping(self._ray_yaml_path,
                                             autostop_config.down)
 


Code review comments:
diff_hunk: @@ -128,7 +128,7 @@ def _run(self):
     def _stop_cluster(self, autostop_config):
         if (autostop_config.backend ==
                 cloud_vm_ray_backend.CloudVmRayBackend.NAME):
-            autostop_lib.set_is_autostopping()
+            autostop_lib.set_autostopping_indicator()
path: sky/skylet/events.py
text: @user1:
I feel "indicator" still suggests it's a binary var/bool. How about `{set_}autostopping_started()`? Not feeling strongly.

@author:
Good point! Renamed. Thanks!