Repository: skypilot-org/skypilot
Pull request: [Serve/Spot] Allow spot queue/cancel/logs during controller INIT state
<!-- Describe the changes in this PR -->

Fixes #1592, #3285

TODO:
- [x] Handle the case, when the controller is actually abnormal (autostopped or manually terminated)

<!-- 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] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
  - ```for i in `seq 1 100`; do sky spot launch -n test-$i --cloud gcp --cpus 2 -y -d "echo hi; sleep 120"; done```; during the submission
  - [x] `sky spot queue`
  - [x] `sky spot logs 10`
  - [x] `sky spot cancel 4 5 6`
- [x] All smoke tests: `pytest tests/test_smoke.py --managed-spot`
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[UX] `sky spot queue/logs` unavailable when there are multiple spot jobs being submitted
From our users:
> The `sky spot queue/logs` will remain unavailable when multiple spot jobs are being submitted.

Two possible ways to solve this:
1. Add an option (or change the meaning of `sky spot queue -r`) so that it waits until the controller is not in INIT mode.
2. Try to run the spot queue query command first when `INIT` status if found instead of directly showing the cache table.

I prefer the second option, as it is more responsive.

Hints from issue discussion:
The second option may lead to some issue that the `skypilot` package is being installed on the cluster.
This is a recurring request from our users. Elevating the priority.

Commit message: address comments

Patch to review:
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 019b2cb4564..87d8cea9f16 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -27,7 +27,7 @@ jobs:
           - tests/test_optimizer_random_dag.py
           - tests/test_storage.py
           - tests/test_wheels.py
-          - tests/test_spot.py
+          - tests/test_spot_serve.py
           - tests/test_yaml_parser.py
     runs-on: ubuntu-latest
     steps:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 9d8c8366f8e..f58492c9391 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -8,6 +8,7 @@
 import re
 import shlex
 import subprocess
+import sys
 import tempfile
 import textwrap
 import time
@@ -1832,7 +1833,6 @@ def run_ray_status_to_check_ray_cluster_healthy() -> bool:
         return record
 
     # All cases below are transitioning the cluster to non-UP states.
-
     if len(node_statuses) > handle.launched_nodes:
         # Unexpected: in the queried region more than 1 cluster with the same
         # constructed name tag returned. This will typically not happen unless
@@ -1953,8 +1953,10 @@ def run_ray_status_to_check_ray_cluster_healthy() -> bool:
 
 
 def _update_cluster_status(
-        cluster_name: str,
-        acquire_per_cluster_status_lock: bool) -> Optional[Dict[str, Any]]:
+    cluster_name: str,
+    acquire_per_cluster_status_lock: bool,
+    cluster_status_lock_timeout: int = CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS
+) -> Optional[Dict[str, Any]]:
     """Update the cluster status.
 
     The cluster status is updated by checking ray cluster and real status from
@@ -1967,9 +1969,9 @@ def _update_cluster_status(
     Args:
         cluster_name: The name of the cluster.
         acquire_per_cluster_status_lock: Whether to acquire the per-cluster lock
-            before updating the status.
-        need_owner_identity_check: Whether to check the owner identity before
-            updating
+          before updating the status.
+        cluster_status_lock_timeout: The timeout to acquire the per-cluster
+          lock.
 
     Returns:
         If the cluster is terminated or does not exist, return None. Otherwise
@@ -1988,23 +1990,22 @@ def _update_cluster_status(
         return _update_cluster_status_no_lock(cluster_name)
 
     try:
-        # TODO(mraheja): remove pylint disabling when filelock
-        # version updated
-        # pylint: disable=abstract-class-instantiated
         with filelock.FileLock(CLUSTER_STATUS_LOCK_PATH.format(cluster_name),
-                               CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS):
+                               timeout=cluster_status_lock_timeout):
             return _update_cluster_status_no_lock(cluster_name)
     except filelock.Timeout:
         logger.debug('Refreshing status: Failed get the lock for cluster '
                      f'{cluster_name!r}. Using the cached status.')
-        return global_user_state.get_cluster_from_name(cluster_name)
+        record = global_user_state.get_cluster_from_name(cluster_name)
+        return record
 
 
 def refresh_cluster_record(
-        cluster_name: str,
-        *,
-        force_refresh_statuses: Optional[Set[status_lib.ClusterStatus]] = None,
-        acquire_per_cluster_status_lock: bool = True
+    cluster_name: str,
+    *,
+    force_refresh_statuses: Optional[Set[status_lib.ClusterStatus]] = None,
+    acquire_per_cluster_status_lock: bool = True,
+    cluster_status_lock_timeout: int = CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS
 ) -> Optional[Dict[str, Any]]:
     """Refresh the cluster, and return the possibly updated record.
 
@@ -2015,13 +2016,15 @@ def refresh_cluster_record(
     Args:
         cluster_name: The name of the cluster.
         force_refresh_statuses: if specified, refresh the cluster if it has one of
-            the specified statuses. Additionally, clusters satisfying the
-            following conditions will always be refreshed no matter the
-            argument is specified or not:
-                1. is a spot cluster, or
-                2. is a non-spot cluster, is not STOPPED, and autostop is set.
+          the specified statuses. Additionally, clusters satisfying the
+          following conditions will always be refreshed no matter the
+          argument is specified or not:
+            1. is a spot cluster, or
+            2. is a non-spot cluster, is not STOPPED, and autostop is set.
         acquire_per_cluster_status_lock: Whether to acquire the per-cluster lock
-            before updating the status.
+          before updating the status.
+        cluster_status_lock_timeout: The timeout to acquire the per-cluster
+          lock. If timeout, the function will use the cached status.
 
     Returns:
         If the cluster is terminated or does not exist, return None.
@@ -2052,7 +2055,8 @@ def refresh_cluster_record(
         if force_refresh_for_cluster or has_autostop or use_spot:
             record = _update_cluster_status(
                 cluster_name,
-                acquire_per_cluster_status_lock=acquire_per_cluster_status_lock)
+                acquire_per_cluster_status_lock=acquire_per_cluster_status_lock,
+                cluster_status_lock_timeout=cluster_status_lock_timeout)
     return record
 
 
@@ -2062,6 +2066,7 @@ def refresh_cluster_status_handle(
     *,
     force_refresh_statuses: Optional[Set[status_lib.ClusterStatus]] = None,
     acquire_per_cluster_status_lock: bool = True,
+    cluster_status_lock_timeout: int = CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS
 ) -> Tuple[Optional[status_lib.ClusterStatus],
            Optional[backends.ResourceHandle]]:
     """Refresh the cluster, and return the possibly updated status and handle.
@@ -2073,7 +2078,8 @@ def refresh_cluster_status_handle(
     record = refresh_cluster_record(
         cluster_name,
         force_refresh_statuses=force_refresh_statuses,
-        acquire_per_cluster_status_lock=acquire_per_cluster_status_lock)
+        acquire_per_cluster_status_lock=acquire_per_cluster_status_lock,
+        cluster_status_lock_timeout=cluster_status_lock_timeout)
     if record is None:
         return None, None
     return record['status'], record['handle']
@@ -2219,41 +2225,48 @@ def check_cluster_available(
 
 
 # TODO(tian): Refactor to controller_utils. Current blocker: circular import.
-def is_controller_up(
+def is_controller_accessible(
     controller_type: controller_utils.Controllers,
     stopped_message: str,
     non_existent_message: Optional[str] = None,
-) -> Tuple[Optional[status_lib.ClusterStatus],
-           Optional['backends.CloudVmRayResourceHandle']]:
+    exit_if_not_accessible: bool = False,
+) -> 'backends.CloudVmRayResourceHandle':
     """Check if the spot/serve controller is up.
 
-    It can be used to check the actual controller status (since the autostop is
-    set for the controller) before the spot/serve commands interact with the
+    The controller is accessible when it is in UP or INIT state, and the ssh
+    connection is successful.
+
+    It can be used to check if the controller is accessible (since the autostop
+    is set for the controller) before the spot/serve commands interact with the
     controller.
 
+    ClusterNotUpError will be raised whenever the controller cannot be accessed.
+
     Args:
         type: Type of the controller.
         stopped_message: Message to print if the controller is STOPPED.
         non_existent_message: Message to show if the controller does not exist.
+        exit_if_not_accessible: Whether to exit directly if the controller is not
+          accessible. If False, the function will raise ClusterNotUpError.
 
     Returns:
-        controller_status: The status of the controller. If it fails during
-          refreshing the status, it will be the cached status. None if the
-          controller does not exist.
-        handle: The ResourceHandle of the controller. None if the
-          controller is not UP or does not exist.
+        handle: The ResourceHandle of the controller.
 
     Raises:
         exceptions.ClusterOwnerIdentityMismatchError: if the current user is not
           the same as the user who created the cluster.
         exceptions.CloudUserIdentityError: if we fail to get the current user
           identity.
+        exceptions.ClusterNotUpError: if the controller is not accessible, or
+          failed to be connected.
     """
     if non_existent_message is None:
         non_existent_message = (
             controller_type.value.default_hint_if_non_existent)
     cluster_name = controller_type.value.cluster_name
     controller_name = controller_type.value.name.replace(' controller', '')
+    need_connection_check = False
+    controller_status, handle = None, None
     try:
         # Set force_refresh_statuses=None to make sure the refresh only happens
         # when the controller is INIT/UP (triggered in these statuses as the
@@ -2261,8 +2274,16 @@ def is_controller_up(
         # unnecessary costly refresh when the controller is already stopped.
         # This optimization is based on the assumption that the user will not
         # start the controller manually from the cloud console.
+        #
+        # The acquire_lock_timeout is set to 0 to avoid hanging the command when
+        # multiple spot_launch commands are running at the same time. Our later
+        # code will check if the controller is accessible by directly checking
+        # the ssh connection to the controller, if it fails to get accurate
+        # status of the controller.
         controller_status, handle = refresh_cluster_status_handle(
-            cluster_name, force_refresh_statuses=None)
+            cluster_name,
+            force_refresh_statuses=None,
+            cluster_status_lock_timeout=0)
     except exceptions.ClusterStatusFetchingError as e:
         # We do not catch the exceptions related to the cluster owner identity
         # mismatch, please refer to the comment in
@@ -2273,22 +2294,51 @@ def is_controller_up(
             'stale information, when the controller is not up.\n'
             f'  Details: {common_utils.format_exception(e, use_bracket=True)}')
         record = global_user_state.get_cluster_from_name(cluster_name)
-        controller_status, handle = None, None
         if record is not None:
             controller_status, handle = record['status'], record['handle']
+            # We check the connection even if the cluster has a cached status UP
+            # to make sure the controller is actually accessible, as the cached
+            # status might be stale.
+            need_connection_check = True
+
+    error_msg = None
+    if controller_status == status_lib.ClusterStatus.STOPPED:
+        error_msg = stopped_message
+    elif controller_status is None or handle is None or handle.head_ip is None:
+        # We check the controller is STOPPED before the check for handle.head_ip
+        # None because when the controller is STOPPED, handle.head_ip can also
+        # be None, but we only want to catch the case when the controller is
+        # being provisioned at the first time and have no head_ip.
+        error_msg = non_existent_message
+    elif (controller_status == status_lib.ClusterStatus.INIT or
+          need_connection_check):
+        # We check the access to controller in INIT state or failed to fetch the
+        # status, as its lock can be hold by another `sky spot launch` or
+        # `sky serve up`, but still have head_ip available. In those cases,
+        # we can allow the access to the controller.
+        ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml,
+                                                   handle.docker_user,
+                                                   handle.ssh_user)
+
+        runner = command_runner.SSHCommandRunner(handle.head_ip,
+                                                 **ssh_credentials,
+                                                 port=handle.head_ssh_port)
+        if not runner.check_connection():
+            error_msg = controller_type.value.connection_error_hint
+    else:
+        assert controller_status == status_lib.ClusterStatus.UP, handle
 
-    if controller_status is None:
-        sky_logging.print(non_existent_message)
-    elif controller_status != status_lib.ClusterStatus.UP:
-        msg = (f'{controller_name.capitalize()} controller {cluster_name} '
-               f'is {controller_status.value}.')
-        if controller_status == status_lib.ClusterStatus.STOPPED:
-            msg += f'\n{stopped_message}'
-        if controller_status == status_lib.ClusterStatus.INIT:
-            msg += '\nPlease wait for the controller to be ready.'
-        sky_logging.print(msg)
-        handle = None
-    return controller_status, handle
+    if error_msg is not None:
+        if exit_if_not_accessible:
+            sky_logging.print(error_msg)
+            sys.exit(1)
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.ClusterNotUpError(error_msg,
+                                               cluster_status=controller_status,
+                                               handle=handle)
+    assert handle is not None and handle.head_ip is not None, (
+        handle, controller_status)
+    return handle
 
 
 class CloudFilter(enum.Enum):
diff --git a/sky/cli.py b/sky/cli.py
index 978bd512a97..45e530cb621 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -1648,15 +1648,10 @@ def _get_spot_jobs(
         num_in_progress_jobs = len(spot_jobs)
     except exceptions.ClusterNotUpError as e:
         controller_status = e.cluster_status
-        if controller_status == status_lib.ClusterStatus.INIT:
-            msg = ('Controller\'s latest status is INIT; jobs '
-                   'will not be shown until it becomes UP.')
-        else:
-            assert controller_status in [None, status_lib.ClusterStatus.STOPPED]
-            msg = 'No in progress jobs.'
-            if controller_status is None:
-                msg += (f' (See: {colorama.Style.BRIGHT}sky spot -h'
-                        f'{colorama.Style.RESET_ALL})')
+        msg = str(e)
+        if controller_status is None:
+            msg += (f' (See: {colorama.Style.BRIGHT}sky spot -h'
+                    f'{colorama.Style.RESET_ALL})')
     except RuntimeError as e:
         msg = ('Failed to query spot jobs due to connection '
                'issues. Try again later. '
@@ -1703,14 +1698,10 @@ def _get_services(service_names: Optional[List[str]],
             num_services = len(service_records)
     except exceptions.ClusterNotUpError as e:
         controller_status = e.cluster_status
-        if controller_status == status_lib.ClusterStatus.INIT:
-            msg = 'Controller is initializing. Please wait for a while.'
-        else:
-            assert controller_status in [None, status_lib.ClusterStatus.STOPPED]
-            msg = 'No existing services. '
-            if controller_status is None:
-                msg += (f'(See: {colorama.Style.BRIGHT}sky serve -h'
-                        f'{colorama.Style.RESET_ALL})')
+        msg = str(e)
+        if controller_status is None:
+            msg += (f'(See: {colorama.Style.BRIGHT}sky serve -h'
+                    f'{colorama.Style.RESET_ALL})')
     except RuntimeError as e:
         msg = ('Failed to fetch service statuses due to connection issues. '
                'Please try again later. Details: '
@@ -2895,86 +2886,74 @@ def down(
 
 
 def _hint_or_raise_for_down_spot_controller(controller_name: str):
-    # spot_jobs will be empty when the spot cluster is not running.
-    cluster_status, _ = backend_utils.refresh_cluster_status_handle(
-        controller_name)
-    if cluster_status is None:
-        click.echo('Managed spot controller has already been torn down.')
-        return
-
     controller = controller_utils.Controllers.from_name(controller_name)
     assert controller is not None, controller_name
-    if cluster_status == status_lib.ClusterStatus.INIT:
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.NotSupportedError(
-                controller.value.decline_down_in_init_status_hint)
+
+    with rich_utils.safe_status(
+            '[bold cyan]Checking for in-progress spot jobs[/]'):
+        try:
+            spot_jobs = core.spot_queue(refresh=False, skip_finished=True)
+        except exceptions.ClusterNotUpError as e:
+            if controller.value.connection_error_hint in str(e):
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.NotSupportedError(
+                        controller.value.
+                        decline_down_when_failed_to_fetch_status_hint)
+            if e.cluster_status is None:
+                click.echo(
+                    'Managed spot controller has already been torn down.')
+                sys.exit(0)
+            # At this point, the spot jobs are failed to be fetched due to the
+            # controller being STOPPED or being firstly launched, i.e., there is
+            # no in-prgress spot jobs.
+            spot_jobs = []
+
     msg = (f'{colorama.Fore.YELLOW}WARNING: Tearing down the managed '
-           f'spot controller ({cluster_status.value}). Please be '
-           f'aware of the following:{colorama.Style.RESET_ALL}'
+           'spot controller. Please be aware of the following:'
+           f'{colorama.Style.RESET_ALL}'
            '\n * All logs and status information of the spot '
            'jobs (output of `sky spot queue`) will be lost.')
     click.echo(msg)
-    if cluster_status == status_lib.ClusterStatus.UP:
-        with rich_utils.safe_status(
-                '[bold cyan]Checking for in-progress spot jobs[/]'):
-            try:
-                spot_jobs = core.spot_queue(refresh=False)
-            except exceptions.ClusterNotUpError:
-                # The spot controller cluster status changed during querying
-                # the spot jobs, use the latest cluster status, so that the
-                # message for INIT and STOPPED states will be correctly
-                # added to the message.
-                cluster_status = backend_utils.refresh_cluster_status_handle(
-                    controller_name)
-                spot_jobs = []
-
-        # Find in-progress spot jobs, and hint users to cancel them.
-        non_terminal_jobs = [
-            job for job in spot_jobs if not job['status'].is_terminal()
-        ]
-        if (cluster_status == status_lib.ClusterStatus.UP and
-                non_terminal_jobs):
-            job_table = spot_lib.format_job_table(non_terminal_jobs,
-                                                  show_all=False)
-            msg = controller.value.decline_down_for_dirty_controller_hint
-            # Add prefix to each line to align with the bullet point.
-            msg += '\n'.join(
-                ['   ' + line for line in job_table.split('\n') if line != ''])
-            with ux_utils.print_exception_no_traceback():
-                raise exceptions.NotSupportedError(msg)
-        else:
-            click.echo(' * No in-progress spot jobs found. It should be safe '
-                       'to terminate (see caveats above).')
+    if spot_jobs:
+        job_table = spot_lib.format_job_table(spot_jobs, show_all=False)
+        msg = controller.value.decline_down_for_dirty_controller_hint
+        # Add prefix to each line to align with the bullet point.
+        msg += '\n'.join(
+            ['   ' + line for line in job_table.split('\n') if line != ''])
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.NotSupportedError(msg)
+    else:
+        click.echo(' * No in-progress spot jobs found. It should be safe to '
+                   'terminate (see caveats above).')
 
 
 def _hint_or_raise_for_down_sky_serve_controller(controller_name: str):
-    cluster_status, _ = backend_utils.refresh_cluster_status_handle(
-        controller_name)
-    if cluster_status is None:
-        click.echo('Sky serve controller has already been torn down.')
-        return
-
     controller = controller_utils.Controllers.from_name(controller_name)
     assert controller is not None, controller_name
-    if cluster_status == status_lib.ClusterStatus.INIT:
+    with rich_utils.safe_status('[bold cyan]Checking for live services[/]'):
+        try:
+            services = serve_lib.status()
+        except exceptions.ClusterNotUpError as e:
+            if controller.value.connection_error_hint in str(e):
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.NotSupportedError(
+                        controller.value.
+                        decline_down_when_failed_to_fetch_status_hint)
+            if e.cluster_status is None:
+                click.echo('Serve controller has already been torn down.')
+                sys.exit(0)
+            # At this point, the services are failed to be fetched due to the
+            # controller being STOPPED or being firstly launched, i.e., there is
+            # no in-prgress services.
+            services = []
+
+    if services:
+        service_names = [service['name'] for service in services]
         with ux_utils.print_exception_no_traceback():
-            raise exceptions.NotSupportedError(
-                controller.value.decline_down_in_init_status_hint)
-    elif cluster_status == status_lib.ClusterStatus.UP:
-        with rich_utils.safe_status(
-                '[bold cyan]Checking for running services[/]'):
-            try:
-                services = serve_lib.status()
-            except exceptions.ClusterNotUpError:
-                cluster_status = backend_utils.refresh_cluster_status_handle(
-                    controller_name)
-                services = []
-        if services:
-            service_names = [service['name'] for service in services]
-            with ux_utils.print_exception_no_traceback():
-                msg = (controller.value.decline_down_for_dirty_controller_hint.
-                       format(service_names=', '.join(service_names)))
-                raise exceptions.NotSupportedError(msg)
+            msg = (
+                controller.value.decline_down_for_dirty_controller_hint.format(
+                    service_names=', '.join(service_names)))
+            raise exceptions.NotSupportedError(msg)
     # Do nothing for STOPPED state, as it is safe to terminate the cluster.
     click.echo(f'Terminate sky serve controller: {controller_name}.')
 
@@ -3064,6 +3043,13 @@ def _down_or_stop_clusters(
                 assert controller is not None
                 hint_or_raise = _CONTROLLER_TO_HINT_OR_RAISE[controller]
                 try:
+                    # TODO(zhwu): This hint or raise is not transactional, which
+                    # means even if it passed the check with no in-progress spot
+                    # or service and prompt the confirmation for termination,
+                    # a user could still do a `sky spot launch` or a
+                    # `sky serve up` before typing the delete, causing a leaked
+                    # spot job or service. We should make this check atomic with
+                    # the termination.
                     hint_or_raise(controller_name)
                 except exceptions.ClusterOwnerIdentityMismatchError as e:
                     if purge:
@@ -4081,12 +4067,10 @@ def spot_cancel(name: Optional[str], job_ids: Tuple[int], all: bool, yes: bool):
       # Cancel managed spot jobs with IDs 1, 2, 3
       $ sky spot cancel 1 2 3
     """
-    _, handle = backend_utils.is_controller_up(
+    backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
-        stopped_message='All managed spot jobs should have finished.')
-    if handle is None:
-        # Hint messages already printed by the call above.
-        sys.exit(1)
+        stopped_message='All managed spot jobs should have finished.',
+        exit_if_not_accessible=True)
 
     job_id_str = ','.join(map(str, job_ids))
     if sum([len(job_ids) > 0, name is not None, all]) != 1:
@@ -4140,8 +4124,8 @@ def spot_logs(name: Optional[str], job_id: Optional[int], follow: bool,
                            follow=follow)
         else:
             core.spot_tail_logs(name=name, job_id=job_id, follow=follow)
-    except exceptions.ClusterNotUpError:
-        # Hint messages already printed by the call above.
+    except exceptions.ClusterNotUpError as e:
+        click.echo(e)
         sys.exit(1)
 
 
@@ -4166,12 +4150,12 @@ def spot_dashboard(port: Optional[int]):
     hint = (
         'Dashboard is not available if spot controller is not up. Run a spot '
         'job first.')
-    _, handle = backend_utils.is_controller_up(
+    backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
         stopped_message=hint,
-        non_existent_message=hint)
-    if handle is None:
-        sys.exit(1)
+        non_existent_message=hint,
+        exit_if_not_accessible=True)
+
     # SSH forward a free local port to remote's dashboard port.
     remote_port = constants.SPOT_DASHBOARD_REMOTE_PORT
     if port is None:
@@ -4674,12 +4658,10 @@ def serve_down(service_names: List[str], all: bool, purge: bool, yes: bool):
             'Can only specify one of SERVICE_NAMES or --all. '
             f'Provided {argument_str!r}.')
 
-    _, handle = backend_utils.is_controller_up(
+    backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
-        stopped_message='All services should have been terminated.')
-    if handle is None:
-        # Hint messages already printed by the call above.
-        sys.exit(1)
+        stopped_message='All services should have been terminated.',
+        exit_if_not_accessible=True)
 
     if not yes:
         quoted_service_names = [f'{name!r}' for name in service_names]
@@ -4759,8 +4741,8 @@ def serve_logs(
                             target=target_component,
                             replica_id=replica_id,
                             follow=follow)
-    except exceptions.ClusterNotUpError:
-        # Hint messages already printed by the call above.
+    except exceptions.ClusterNotUpError as e:
+        click.echo(e)
         sys.exit(1)
 
 
diff --git a/sky/core.py b/sky/core.py
index 10e300ca594..0639d14d206 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -796,17 +796,21 @@ def spot_queue(refresh: bool,
             does not exist.
         RuntimeError: if failed to get the spot jobs with ssh.
     """
-
     stop_msg = ''
     if not refresh:
-        stop_msg = 'To view the latest job table: sky spot queue --refresh'
-    controller_status, handle = backend_utils.is_controller_up(
-        controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
-        stopped_message=stop_msg)
+        stop_msg = ('No in-progress spot jobs.\nTo view the latest job table: '
+                    'sky spot queue --refresh')
+    try:
+        handle = backend_utils.is_controller_accessible(
+            controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
+            stopped_message=stop_msg)
+    except exceptions.ClusterNotUpError as e:
+        if not refresh:
+            raise
+        handle = None
+        controller_status = e.cluster_status
 
-    if (refresh and controller_status in [
-            status_lib.ClusterStatus.STOPPED, status_lib.ClusterStatus.INIT
-    ]):
+    if refresh and handle is None:
         sky_logging.print(f'{colorama.Fore.YELLOW}'
                           'Restarting controller for latest status...'
                           f'{colorama.Style.RESET_ALL}')
@@ -817,15 +821,7 @@ def spot_queue(refresh: bool,
         controller_status = status_lib.ClusterStatus.UP
         rich_utils.force_update_status('[cyan] Checking spot jobs[/]')
 
-    if handle is None or handle.head_ip is None:
-        # When the controller is STOPPED, the head_ip will be None, as
-        # it will be set in global_user_state.remove_cluster().
-        # We do not directly check for UP because the controller may be
-        # in INIT state during another spot launch, but still have
-        # head_ip available. In this case, we can still try to ssh
-        # into the controller and fetch the job table.
-        raise exceptions.ClusterNotUpError('Spot controller is not up.',
-                                           cluster_status=controller_status)
+    assert handle is not None, (controller_status, refresh)
 
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
@@ -837,14 +833,15 @@ def spot_queue(refresh: bool,
         require_outputs=True,
         stream_logs=False,
         separate_stderr=True)
+
     try:
         subprocess_utils.handle_returncode(returncode,
                                            code,
                                            'Failed to fetch managed spot jobs',
-                                           stderr,
+                                           job_table_payload + stderr,
                                            stream_logs=False)
     except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+        raise RuntimeError(str(e)) from e
 
     jobs = spot.load_spot_job_queue(job_table_payload)
     if skip_finished:
@@ -873,16 +870,9 @@ def spot_cancel(name: Optional[str] = None,
         RuntimeError: failed to cancel the job.
     """
     job_ids = [] if job_ids is None else job_ids
-    cluster_status, handle = backend_utils.is_controller_up(
+    handle = backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
         stopped_message='All managed spot jobs should have finished.')
-    if handle is None or handle.head_ip is None:
-        # The error message is already printed in
-        # backend_utils.is_controller_up
-        # TODO(zhwu): Move the error message into the exception.
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ClusterNotUpError(message='',
-                                               cluster_status=cluster_status)
 
     job_id_str = ','.join(map(str, job_ids))
     if sum([len(job_ids) > 0, name is not None, all]) != 1:
@@ -933,17 +923,10 @@ def spot_tail_logs(name: Optional[str], job_id: Optional[int],
         sky.exceptions.ClusterNotUpError: the spot controller is not up.
     """
     # TODO(zhwu): Automatically restart the spot controller
-    controller_status, handle = backend_utils.is_controller_up(
+    handle = backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
         stopped_message=('Please restart the spot controller with '
                          f'`sky start {spot.SPOT_CONTROLLER_NAME}`.'))
-    if handle is None or handle.head_ip is None:
-        msg = 'All jobs finished.'
-        if controller_status == status_lib.ClusterStatus.INIT:
-            msg = ''
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ClusterNotUpError(msg,
-                                               cluster_status=controller_status)
 
     if name is not None and job_id is not None:
         raise ValueError('Cannot specify both name and job_id.')
diff --git a/sky/serve/core.py b/sky/serve/core.py
index cba2f05b8d0..e9a41761bfd 100644
--- a/sky/serve/core.py
+++ b/sky/serve/core.py
@@ -1,7 +1,6 @@
 """SkyServe core APIs."""
 import re
 import tempfile
-import typing
 from typing import Any, Dict, List, Optional, Union
 
 import colorama
@@ -11,7 +10,6 @@
 from sky import exceptions
 from sky import global_user_state
 from sky import sky_logging
-from sky import status_lib
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.serve import constants as serve_constants
@@ -26,9 +24,6 @@
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
-if typing.TYPE_CHECKING:
-    from sky import clouds
-
 logger = sky_logging.init_logger(__name__)
 
 
@@ -310,7 +305,7 @@ def update(task: 'sky.Task', service_name: str) -> None:
         service_name: Name of the service.
     """
     _validate_service_task(task)
-    cluster_status, handle = backend_utils.is_controller_up(
+    handle = backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
         stopped_message=
         'Service controller is stopped. There is no service to update. '
@@ -321,14 +316,6 @@ def update(task: 'sky.Task', service_name: str) -> None:
         f'use {backend_utils.BOLD}sky serve up{backend_utils.RESET_BOLD}',
     )
 
-    if handle is None or handle.head_ip is None:
-        # The error message is already printed in
-        # backend_utils.is_controller_up
-        # TODO(zhwu): Move the error message into the exception.
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ClusterNotUpError(message='',
-                                               cluster_status=cluster_status)
-
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
 
@@ -461,16 +448,9 @@ def down(
         service_names = []
     if isinstance(service_names, str):
         service_names = [service_names]
-    cluster_status, handle = backend_utils.is_controller_up(
+    handle = backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
         stopped_message='All services should have terminated.')
-    if handle is None or handle.head_ip is None:
-        # The error message is already printed in
-        # backend_utils.is_controller_up
-        # TODO(zhwu): Move the error message into the exception.
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ClusterNotUpError(message='',
-                                               cluster_status=cluster_status)
 
     service_names_str = ','.join(service_names)
     if sum([len(service_names) > 0, all]) != 1:
@@ -573,21 +553,10 @@ def status(
             raise RuntimeError(
                 'Failed to refresh service status due to network error.') from e
 
-    # TODO(tian): This is so slow... It will take ~10s to refresh the status
-    # of controller. Can we optimize this?
-    controller_status, handle = backend_utils.is_controller_up(
-        controller_type=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
-        stopped_message='No service is found.')
-
-    if handle is None or handle.head_ip is None:
-        # When the controller is STOPPED, the head_ip will be None, as
-        # it will be set in global_user_state.remove_cluster().
-        # We do not directly check for UP because the controller may be
-        # in INIT state during another `sky serve up`, but still have
-        # head_ip available. In this case, we can still try to ssh
-        # into the controller and fetch the job table.
-        raise exceptions.ClusterNotUpError('Sky serve controller is not up.',
-                                           cluster_status=controller_status)
+    controller_type = controller_utils.Controllers.SKY_SERVE_CONTROLLER
+    handle = backend_utils.is_controller_accessible(
+        controller_type=controller_type,
+        stopped_message=controller_type.value.default_hint_if_non_existent)
 
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
@@ -669,15 +638,10 @@ def tail_logs(
             with ux_utils.print_exception_no_traceback():
                 raise ValueError('`replica_id` must be None when using '
                                  'target=CONTROLLER/LOAD_BALANCER.')
-    controller_status, handle = backend_utils.is_controller_up(
+    handle = backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
-        stopped_message='No service is found.')
-    if handle is None or handle.head_ip is None:
-        msg = 'No service is found.'
-        if controller_status == status_lib.ClusterStatus.INIT:
-            msg = ''
-        raise exceptions.ClusterNotUpError(msg,
-                                           cluster_status=controller_status)
+        stopped_message='No live services.')
+
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend), backend
     backend.tail_serve_logs(handle,
diff --git a/sky/sky_logging.py b/sky/sky_logging.py
index 2a590f950ee..dbaf1dd0479 100644
--- a/sky/sky_logging.py
+++ b/sky/sky_logging.py
@@ -111,12 +111,13 @@ def silent():
     _root_logger.setLevel(logging.ERROR)
     _logging_config.is_silent = True
     print = lambda *args, **kwargs: None
-    yield
-
-    # Restore logger
-    print = previous_print
-    _root_logger.setLevel(previous_level)
-    _logging_config.is_silent = previous_is_silent
+    try:
+        yield
+    finally:
+        # Restore logger
+        print = previous_print
+        _root_logger.setLevel(previous_level)
+        _logging_config.is_silent = previous_is_silent
 
 
 def is_silent():
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 815cc52c246..e4715497f39 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -657,7 +657,7 @@ def format_job_table(
     if status_str:
         status_str = f'In progress tasks: {status_str}'
     else:
-        status_str = 'No in progress tasks.'
+        status_str = 'No in-progress spot jobs.'
     output = status_str
     if str(job_table):
         output += f'\n{job_table}'
diff --git a/sky/utils/command_runner.py b/sky/utils/command_runner.py
index f7ec5bd4f41..b8036d50bab 100644
--- a/sky/utils/command_runner.py
+++ b/sky/utils/command_runner.py
@@ -47,11 +47,13 @@ def ssh_options_list(
     *,
     ssh_proxy_command: Optional[str] = None,
     docker_ssh_proxy_command: Optional[str] = None,
-    timeout: int = 30,
+    connect_timeout: Optional[int] = None,
     port: int = 22,
     disable_control_master: Optional[bool] = False,
 ) -> List[str]:
     """Returns a list of sane options for 'ssh'."""
+    if connect_timeout is None:
+        connect_timeout = 30
     # Forked from Ray SSHOptions:
     # https://github.com/ray-project/ray/blob/master/python/ray/autoscaler/_private/command_runner.py
     arg_dict = {
@@ -75,7 +77,7 @@ def ssh_options_list(
         'ServerAliveInterval': 5,
         'ServerAliveCountMax': 3,
         # ConnectTimeout.
-        'ConnectTimeout': f'{timeout}s',
+        'ConnectTimeout': f'{connect_timeout}s',
         # Agent forwarding for git.
         'ForwardAgent': 'yes',
     }
@@ -218,7 +220,8 @@ def make_runner_list(
         ]
 
     def _ssh_base_command(self, *, ssh_mode: SshMode,
-                          port_forward: Optional[List[int]]) -> List[str]:
+                          port_forward: Optional[List[int]],
+                          connect_timeout: Optional[int]) -> List[str]:
         ssh = ['ssh']
         if ssh_mode == SshMode.NON_INTERACTIVE:
             # Disable pseudo-terminal allocation. Otherwise, the output of
@@ -243,6 +246,7 @@ def _ssh_base_command(self, *, ssh_mode: SshMode,
             ssh_proxy_command=self._ssh_proxy_command,
             docker_ssh_proxy_command=docker_ssh_proxy_command,
             port=self.port,
+            connect_timeout=connect_timeout,
             disable_control_master=self.disable_control_master) + [
                 f'{self.ssh_user}@{self.ip}'
             ]
@@ -260,6 +264,7 @@ def run(
             stream_logs: bool = True,
             ssh_mode: SshMode = SshMode.NON_INTERACTIVE,
             separate_stderr: bool = False,
+            connect_timeout: Optional[int] = None,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Uses 'ssh' to run 'cmd' on a node with ip.
 
@@ -285,8 +290,10 @@ def run(
             or
             A tuple of (returncode, stdout, stderr).
         """
-        base_ssh_command = self._ssh_base_command(ssh_mode=ssh_mode,
-                                                  port_forward=port_forward)
+        base_ssh_command = self._ssh_base_command(
+            ssh_mode=ssh_mode,
+            port_forward=port_forward,
+            connect_timeout=connect_timeout)
         if ssh_mode == SshMode.LOGIN:
             assert isinstance(cmd, list), 'cmd must be a list for login mode.'
             command = base_ssh_command + cmd
@@ -449,3 +456,10 @@ def rsync(
                                            error_msg,
                                            stderr=stderr,
                                            stream_logs=stream_logs)
+
+    def check_connection(self) -> bool:
+        """Check if the connection to the remote machine is successful."""
+        returncode = self.run('true', connect_timeout=5)
+        if returncode:
+            return False
+        return True
diff --git a/sky/utils/command_runner.pyi b/sky/utils/command_runner.pyi
index 893acb1e57b..f1b547927d3 100644
--- a/sky/utils/command_runner.pyi
+++ b/sky/utils/command_runner.pyi
@@ -123,3 +123,6 @@ class SSHCommandRunner:
               log_path: str = ...,
               stream_logs: bool = ...) -> None:
         ...
+
+    def check_connection(self) -> bool:
+        ...
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
index 6d9ecf04cce..08c25d8100c 100644
--- a/sky/utils/controller_utils.py
+++ b/sky/utils/controller_utils.py
@@ -51,10 +51,11 @@ class _ControllerSpec:
     cluster_name: str
     in_progress_hint: str
     decline_cancel_hint: str
-    decline_down_in_init_status_hint: str
+    decline_down_when_failed_to_fetch_status_hint: str
     decline_down_for_dirty_controller_hint: str
     check_cluster_name_hint: str
     default_hint_if_non_existent: str
+    connection_error_hint: str
 
 
 class Controllers(enum.Enum):
@@ -71,7 +72,7 @@ class Controllers(enum.Enum):
             'Cancelling the spot controller\'s jobs is not allowed.\nTo cancel '
             f'spot jobs, use: {colorama.Style.BRIGHT}sky spot cancel <spot '
             f'job IDs> [--all]{colorama.Style.RESET_ALL}'),
-        decline_down_in_init_status_hint=(
+        decline_down_when_failed_to_fetch_status_hint=(
             f'{colorama.Fore.RED}Tearing down the spot controller while '
             'it is in INIT state is not supported (this means a spot launch '
             'is in progress or the previous launch failed), as we cannot '
@@ -85,17 +86,19 @@ class Controllers(enum.Enum):
             f'sky spot cancel -a{colorama.Style.RESET_ALL}\n'),
         check_cluster_name_hint=(
             f'Cluster {spot_utils.SPOT_CONTROLLER_NAME} is reserved for '
-            'managed spot controller. '),
-        default_hint_if_non_existent='No managed spot jobs are found.')
+            'managed spot controller.'),
+        default_hint_if_non_existent='No in-progress spot jobs.',
+        connection_error_hint=(
+            'Failed to connect to spot controller, please try again later.'))
     SKY_SERVE_CONTROLLER = _ControllerSpec(
-        name='sky serve controller',
+        name='serve controller',
         cluster_name=serve_utils.SKY_SERVE_CONTROLLER_NAME,
         in_progress_hint=(
             f'* To see detailed service status: {colorama.Style.BRIGHT}'
             f'sky serve status -a{colorama.Style.RESET_ALL}'),
         decline_cancel_hint=(
             'Cancelling the sky serve controller\'s jobs is not allowed.'),
-        decline_down_in_init_status_hint=(
+        decline_down_when_failed_to_fetch_status_hint=(
             f'{colorama.Fore.RED}Tearing down the sky serve controller '
             'while it is in INIT state is not supported (this means a sky '
             'serve up is in progress or the previous launch failed), as we '
@@ -112,8 +115,10 @@ class Controllers(enum.Enum):
             f'{colorama.Style.RESET_ALL}.'),
         check_cluster_name_hint=(
             f'Cluster {serve_utils.SKY_SERVE_CONTROLLER_NAME} is reserved for '
-            'sky serve controller. '),
-        default_hint_if_non_existent='No service is found.')
+            'sky serve controller.'),
+        default_hint_if_non_existent='No live services.',
+        connection_error_hint=(
+            'Failed to connect to serve controller, please try again later.'))
 
     @classmethod
     def from_name(cls, name: Optional[str]) -> Optional['Controllers']:
diff --git a/tests/test_spot.py b/tests/test_spot.py
deleted file mode 100644
index 9c7b0e3df81..00000000000
--- a/tests/test_spot.py
+++ /dev/null
@@ -1,184 +0,0 @@
-import tempfile
-import textwrap
-
-import click
-from click import testing as cli_testing
-import pytest
-
-import sky
-from sky import backends
-from sky import cli
-from sky import global_user_state
-from sky import spot
-from sky.utils import db_utils
-
-
-def test_spot_nonexist_strategy():
-    """Test the nonexist recovery strategy."""
-    task_yaml = textwrap.dedent("""\
-        resources:
-            cloud: aws
-            use_spot: true
-            spot_recovery: nonexist""")
-    with tempfile.NamedTemporaryFile(mode='w') as f:
-        f.write(task_yaml)
-        f.flush()
-        with pytest.raises(
-                ValueError,
-                match='is not supported. The strategy should be among'):
-            sky.Task.from_yaml(f.name)
-
-
-class TestControllerOperations:
-    """Test operations on controllers."""
-
-    @pytest.fixture
-    def _mock_db_conn(self, monkeypatch, tmp_path):
-        tmp_path.mkdir(parents=True, exist_ok=True)
-        db_path = tmp_path / 'state_testing.db'
-        monkeypatch.setattr(
-            global_user_state, '_DB',
-            db_utils.SQLiteConn(str(db_path), global_user_state.create_table))
-
-    @pytest.fixture
-    def _mock_cluster_state(self, _mock_db_conn):
-        assert 'state.db' not in global_user_state._DB.db_path
-        handle = backends.CloudVmRayResourceHandle(
-            cluster_name='test-cluster1',
-            cluster_name_on_cloud='test-cluster1',
-            cluster_yaml='/tmp/cluster1.yaml',
-            launched_nodes=2,
-            launched_resources=sky.Resources(sky.AWS(),
-                                             instance_type='p3.2xlarge',
-                                             region='us-east-1'),
-        )
-        global_user_state.add_or_update_cluster(
-            'test-cluster1',
-            handle,
-            requested_resources={handle.launched_resources},
-            ready=True)
-        handle = backends.CloudVmRayResourceHandle(
-            cluster_name='test-cluster2',
-            cluster_name_on_cloud='test-cluster2',
-            cluster_yaml='/tmp/cluster2.yaml',
-            launched_nodes=1,
-            launched_resources=sky.Resources(sky.GCP(),
-                                             instance_type='a2-highgpu-4g',
-                                             accelerators={'A100': 4},
-                                             region='us-west1'),
-        )
-        global_user_state.add_or_update_cluster(
-            'test-cluster2',
-            handle,
-            requested_resources={handle.launched_resources},
-            ready=True)
-        handle = backends.CloudVmRayResourceHandle(
-            cluster_name='test-cluster3',
-            cluster_name_on_cloud='test-cluster3',
-            cluster_yaml='/tmp/cluster3.yaml',
-            launched_nodes=4,
-            launched_resources=sky.Resources(sky.Azure(),
-                                             instance_type='Standard_D4s_v3',
-                                             region='eastus'),
-        )
-        global_user_state.add_or_update_cluster(
-            'test-cluster3',
-            handle,
-            requested_resources={handle.launched_resources},
-            ready=False)
-        handle = backends.CloudVmRayResourceHandle(
-            cluster_name=spot.SPOT_CONTROLLER_NAME,
-            cluster_name_on_cloud=spot.SPOT_CONTROLLER_NAME,
-            cluster_yaml='/tmp/spot_controller.yaml',
-            launched_nodes=1,
-            launched_resources=sky.Resources(sky.AWS(),
-                                             instance_type='m4.2xlarge',
-                                             region='us-west-1'),
-        )
-        global_user_state.add_or_update_cluster(
-            spot.SPOT_CONTROLLER_NAME,
-            handle,
-            requested_resources={handle.launched_resources},
-            ready=True)
-
-    @pytest.mark.timeout(60)
-    def test_down_spot_controller(self, _mock_cluster_state, monkeypatch):
-
-        def mock_cluster_refresh_up(
-            cluster_name: str,
-            *,
-            force_refresh_statuses: bool = False,
-            acquire_per_cluster_status_lock: bool = True,
-        ):
-            record = global_user_state.get_cluster_from_name(cluster_name)
-            return record['status'], record['handle']
-
-        monkeypatch.setattr(
-            'sky.backends.backend_utils.refresh_cluster_status_handle',
-            mock_cluster_refresh_up)
-
-        monkeypatch.setattr('sky.core.spot_queue', lambda refresh: [])
-
-        cli_runner = cli_testing.CliRunner()
-        result = cli_runner.invoke(cli.down, [spot.SPOT_CONTROLLER_NAME],
-                                   input='n')
-        assert 'WARNING: Tearing down the managed spot controller (UP).' in result.output
-        assert isinstance(result.exception,
-                          SystemExit), (result.exception, result.output)
-
-        result = cli_runner.invoke(cli.down, ['sky-spot-con*'])
-        assert not result.exception
-        assert 'Cluster(s) not found' in result.output
-
-        result = cli_runner.invoke(cli.down, ['sky-spot-con*', '-p'])
-        assert not result.exception
-        assert 'Cluster(s) not found' in result.output
-
-        result = cli_runner.invoke(cli.down, ['-a'], input='n\n')
-        assert isinstance(result.exception, SystemExit), result.exception
-        assert 'Aborted' in result.output
-
-        result = cli_runner.invoke(cli.down, ['-ap'], input='n\n')
-        assert isinstance(result.exception, SystemExit)
-        assert 'Aborted' in result.output
-
-    @pytest.mark.timeout(60)
-    def test_stop_spot_controller(self, _mock_cluster_state):
-        cli_runner = cli_testing.CliRunner()
-        result = cli_runner.invoke(cli.stop, [spot.SPOT_CONTROLLER_NAME])
-        assert result.exit_code == click.UsageError.exit_code
-        assert (f'Stopping controller(s) \'{spot.SPOT_CONTROLLER_NAME}\' is '
-                'currently not supported' in result.output)
-
-        result = cli_runner.invoke(cli.stop, ['sky-spot-con*'])
-        assert not result.exception
-        assert 'Cluster(s) not found' in result.output
-
-        result = cli_runner.invoke(cli.stop, ['-a'], input='n\n')
-        assert isinstance(result.exception, SystemExit)
-        assert 'Aborted' in result.output
-
-    @pytest.mark.timeout(60)
-    def test_autostop_spot_controller(self, _mock_cluster_state):
-        cli_runner = cli_testing.CliRunner()
-        result = cli_runner.invoke(cli.autostop, [spot.SPOT_CONTROLLER_NAME])
-        assert result.exit_code == click.UsageError.exit_code
-        assert ('Scheduling autostop on controller(s) '
-                f'\'{spot.SPOT_CONTROLLER_NAME}\' is currently not supported'
-                in result.output)
-
-        result = cli_runner.invoke(cli.autostop, ['sky-spot-con*'])
-        assert not result.exception
-        assert 'Cluster(s) not found' in result.output
-
-        result = cli_runner.invoke(cli.autostop, ['-a'], input='n\n')
-        assert isinstance(result.exception, SystemExit)
-        assert 'Aborted' in result.output
-
-    def test_cancel_on_spot_controller(self, _mock_cluster_state):
-        cli_runner = cli_testing.CliRunner()
-        result = cli_runner.invoke(cli.cancel,
-                                   [spot.SPOT_CONTROLLER_NAME, '-a'])
-        assert result.exit_code == 1
-        assert 'Cancelling the spot controller\'s jobs is not allowed.' in str(
-            result.output)
diff --git a/tests/test_spot_serve.py b/tests/test_spot_serve.py
new file mode 100644
index 00000000000..15dcde6a8de
--- /dev/null
+++ b/tests/test_spot_serve.py
@@ -0,0 +1,428 @@
+import base64
+import pickle
+import tempfile
+import textwrap
+import time
+from typing import Optional, Tuple
+
+import click
+from click import testing as cli_testing
+import pytest
+
+import sky
+from sky import backends
+from sky import cli
+from sky import exceptions
+from sky import global_user_state
+from sky import serve
+from sky import spot
+from sky.utils import common_utils
+from sky.utils import controller_utils
+from sky.utils import db_utils
+
+
+def test_spot_nonexist_strategy():
+    """Test the nonexist recovery strategy."""
+    task_yaml = textwrap.dedent("""\
+        resources:
+            cloud: aws
+            use_spot: true
+            spot_recovery: nonexist""")
+    with tempfile.NamedTemporaryFile(mode='w') as f:
+        f.write(task_yaml)
+        f.flush()
+        with pytest.raises(
+                ValueError,
+                match='is not supported. The strategy should be among'):
+            sky.Task.from_yaml(f.name)
+
+
+@pytest.fixture
+def _mock_db_conn(monkeypatch, tmp_path):
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    db_path = tmp_path / 'state_testing.db'
+    monkeypatch.setattr(
+        global_user_state, '_DB',
+        db_utils.SQLiteConn(str(db_path), global_user_state.create_table))
+
+
+@pytest.fixture
+def _mock_cluster_state(_mock_db_conn):
+    assert 'state.db' not in global_user_state._DB.db_path
+    handle = backends.CloudVmRayResourceHandle(
+        cluster_name='test-cluster1',
+        cluster_name_on_cloud='test-cluster1',
+        cluster_yaml='/tmp/cluster1.yaml',
+        launched_nodes=2,
+        launched_resources=sky.Resources(sky.AWS(),
+                                         instance_type='p3.2xlarge',
+                                         region='us-east-1'),
+    )
+    global_user_state.add_or_update_cluster(
+        'test-cluster1',
+        handle,
+        requested_resources={handle.launched_resources},
+        ready=True)
+    handle = backends.CloudVmRayResourceHandle(
+        cluster_name='test-cluster2',
+        cluster_name_on_cloud='test-cluster2',
+        cluster_yaml='/tmp/cluster2.yaml',
+        launched_nodes=1,
+        launched_resources=sky.Resources(sky.GCP(),
+                                         instance_type='a2-highgpu-4g',
+                                         accelerators={'A100': 4},
+                                         region='us-west1'),
+    )
+    global_user_state.add_or_update_cluster(
+        'test-cluster2',
+        handle,
+        requested_resources={handle.launched_resources},
+        ready=True)
+    handle = backends.CloudVmRayResourceHandle(
+        cluster_name='test-cluster3',
+        cluster_name_on_cloud='test-cluster3',
+        cluster_yaml='/tmp/cluster3.yaml',
+        launched_nodes=4,
+        launched_resources=sky.Resources(sky.Azure(),
+                                         instance_type='Standard_D4s_v3',
+                                         region='eastus'),
+    )
+    global_user_state.add_or_update_cluster(
+        'test-cluster3',
+        handle,
+        requested_resources={handle.launched_resources},
+        ready=False)
+
+
+@pytest.fixture
+def _mock_spot_controller(_mock_db_conn):
+    handle = backends.CloudVmRayResourceHandle(
+        cluster_name=spot.SPOT_CONTROLLER_NAME,
+        cluster_name_on_cloud=spot.SPOT_CONTROLLER_NAME,
+        cluster_yaml='/tmp/spot_controller.yaml',
+        launched_nodes=1,
+        launched_resources=sky.Resources(sky.AWS(),
+                                         instance_type='m4.2xlarge',
+                                         region='us-west-1'),
+    )
+    global_user_state.add_or_update_cluster(
+        spot.SPOT_CONTROLLER_NAME,
+        handle,
+        requested_resources={handle.launched_resources},
+        ready=True)
+
+
+@pytest.fixture
+def _mock_serve_controller(_mock_db_conn):
+    handle = backends.CloudVmRayResourceHandle(
+        cluster_name=serve.SKY_SERVE_CONTROLLER_NAME,
+        cluster_name_on_cloud=serve.SKY_SERVE_CONTROLLER_NAME,
+        cluster_yaml='/tmp/serve_controller.yaml',
+        launched_nodes=1,
+        launched_resources=sky.Resources(sky.AWS(),
+                                         instance_type='m4.2xlarge',
+                                         region='us-west-1'),
+    )
+    global_user_state.add_or_update_cluster(
+        serve.SKY_SERVE_CONTROLLER_NAME,
+        handle,
+        requested_resources={handle.launched_resources},
+        ready=True)
+
+
+def mock_is_controller_accessible(
+    controller_type: controller_utils.Controllers,
+    stopped_message: str,
+    non_existent_message: Optional[str] = None,
+    exit_on_error: bool = False,
+):
+    record = global_user_state.get_cluster_from_name(
+        controller_type.value.cluster_name)
+    return record['handle']
+
+
+class TestSpotOperations:
+    """Test operations for managed spot."""
+
+    @pytest.mark.timeout(60)
+    def test_down_spot_controller(self, _mock_cluster_state,
+                                  _mock_spot_controller, monkeypatch):
+
+        def mock_get_job_table_no_job(cls, handle, code, require_outputs,
+                                      stream_logs,
+                                      separate_stderr) -> Tuple[int, str, str]:
+            return 0, common_utils.encode_payload([]), ''
+
+        def mock_get_job_table_one_job(cls, handle, code, require_outputs,
+                                       stream_logs,
+                                       separate_stderr) -> Tuple[int, str, str]:
+            return 0, common_utils.encode_payload([{
+                'job_id': '1',
+                'job_name': 'test_job',
+                'resources': 'test',
+                'status': 'RUNNING',
+                'submitted_at': time.time(),
+                'run_timestamp': str(time.time()),
+                'start_at': time.time(),
+                'end_at': time.time(),
+                'last_recovered_at': None,
+                'recovery_count': 0,
+                'failure_reason': '',
+                'spot_job_id': '1',
+                'task_id': 0,
+                'task_name': 'test_task',
+                'job_duration': 20,
+            }]), ''
+
+        monkeypatch.setattr(
+            'sky.backends.backend_utils.is_controller_accessible',
+            mock_is_controller_accessible)
+
+        monkeypatch.setattr(
+            'sky.backends.cloud_vm_ray_backend.CloudVmRayBackend.run_on_head',
+            mock_get_job_table_no_job)
+
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.down, [spot.SPOT_CONTROLLER_NAME],
+                                   input='n')
+        assert 'WARNING: Tearing down the managed spot controller.' in result.output, (
+            result.exception, result.output, result.exc_info)
+        assert isinstance(result.exception,
+                          SystemExit), (result.exception, result.output)
+
+        monkeypatch.setattr(
+            'sky.backends.cloud_vm_ray_backend.CloudVmRayBackend.run_on_head',
+            mock_get_job_table_one_job)
+        result = cli_runner.invoke(cli.down, [spot.SPOT_CONTROLLER_NAME],
+                                   input='n')
+        assert 'WARNING: Tearing down the managed spot controller.' in result.output, (
+            result.exception, result.output, result.exc_info)
+        assert isinstance(result.exception, exceptions.NotSupportedError)
+
+        result = cli_runner.invoke(cli.down, ['sky-spot-con*'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.down, ['sky-spot-con*', '-p'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.down, ['-a'], input='n\n')
+        assert isinstance(result.exception, SystemExit), result.exception
+        assert 'Aborted' in result.output
+
+        result = cli_runner.invoke(cli.down, ['-ap'], input='n\n')
+        assert isinstance(result.exception, SystemExit)
+        assert 'Aborted' in result.output
+
+    @pytest.mark.timeout(60)
+    def test_stop_spot_controller(self, _mock_cluster_state,
+                                  _mock_spot_controller):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.stop, [spot.SPOT_CONTROLLER_NAME])
+        assert result.exit_code == click.UsageError.exit_code
+        assert (f'Stopping controller(s) \'{spot.SPOT_CONTROLLER_NAME}\' is '
+                'currently not supported' in result.output)
+
+        result = cli_runner.invoke(cli.stop, ['sky-spot-con*'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.stop, ['-a'], input='n\n')
+        assert isinstance(result.exception, SystemExit)
+        assert 'Aborted' in result.output
+
+    @pytest.mark.timeout(60)
+    def test_autostop_spot_controller(self, _mock_cluster_state,
+                                      _mock_spot_controller):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.autostop, [spot.SPOT_CONTROLLER_NAME])
+        assert result.exit_code == click.UsageError.exit_code
+        assert ('Scheduling autostop on controller(s) '
+                f'\'{spot.SPOT_CONTROLLER_NAME}\' is currently not supported'
+                in result.output)
+
+        result = cli_runner.invoke(cli.autostop, ['sky-spot-con*'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.autostop, ['-a'], input='n\n')
+        assert isinstance(result.exception, SystemExit)
+        assert 'Aborted' in result.output
+
+    def test_cancel_on_spot_controller(self, _mock_cluster_state,
+                                       _mock_spot_controller):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.cancel,
+                                   [spot.SPOT_CONTROLLER_NAME, '-a'])
+        assert result.exit_code == 1
+        assert 'Cancelling the spot controller\'s jobs is not allowed.' in str(
+            result.output)
+
+    @pytest.mark.timeout(60)
+    def test_cancel(self, _mock_db_conn):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.spot_cancel, ['-a'])
+        assert result.exit_code == 1
+        assert controller_utils.Controllers.SPOT_CONTROLLER.value.default_hint_if_non_existent in str(
+            result.output), (result.exception, result.output, result.exc_info)
+
+    @pytest.mark.timeout(60)
+    def test_logs(self, _mock_db_conn):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.spot_logs, ['1'])
+        assert result.exit_code == 1
+        assert controller_utils.Controllers.SPOT_CONTROLLER.value.default_hint_if_non_existent in str(
+            result.output), (result.exception, result.output, result.exc_info)
+
+    @pytest.mark.timeout(60)
+    def test_queue(self, _mock_db_conn):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.spot_queue)
+        assert result.exit_code == 0
+        assert controller_utils.Controllers.SPOT_CONTROLLER.value.default_hint_if_non_existent in str(
+            result.output), (result.exception, result.output, result.exc_info)
+
+
+class TestServeOperations:
+    """Test operations for services."""
+
+    @pytest.mark.timeout(60)
+    def test_down_spot_controller(self, _mock_cluster_state,
+                                  _mock_serve_controller, monkeypatch):
+
+        def mock_get_services_no_service(
+                cls, handle, code, require_outputs, stream_logs,
+                separate_stderr) -> Tuple[int, str, str]:
+            return 0, common_utils.encode_payload([]), ''
+
+        def mock_get_services_one_service(
+                cls, handle, code, require_outputs, stream_logs,
+                separate_stderr) -> Tuple[int, str, str]:
+            service = {
+                'name': 'test_service',
+                'controller_job_id': 1,
+                'uptime': 20,
+                'status': 'RUNNING',
+                'controller_port': 30001,
+                'load_balancer_port': 30000,
+                'policy': None,
+                'requested_resources': sky.Resources(),
+                'requested_resources_str': '',
+                'replica_info': [],
+            }
+            return 0, common_utils.encode_payload([{
+                k: base64.b64encode(pickle.dumps(v)).decode('utf-8')
+                for k, v in service.items()
+            }]), ''
+
+        monkeypatch.setattr(
+            'sky.backends.backend_utils.is_controller_accessible',
+            mock_is_controller_accessible)
+
+        monkeypatch.setattr(
+            'sky.backends.cloud_vm_ray_backend.CloudVmRayBackend.run_on_head',
+            mock_get_services_no_service)
+
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.down, [serve.SKY_SERVE_CONTROLLER_NAME],
+                                   input='n')
+        assert 'Terminate sky serve controller:' in result.output, (
+            result.exception, result.output, result.exc_info)
+        assert isinstance(result.exception,
+                          SystemExit), (result.exception, result.output)
+
+        monkeypatch.setattr(
+            'sky.backends.cloud_vm_ray_backend.CloudVmRayBackend.run_on_head',
+            mock_get_services_one_service)
+        result = cli_runner.invoke(cli.down, [serve.SKY_SERVE_CONTROLLER_NAME],
+                                   input='n')
+        assert isinstance(result.exception, exceptions.NotSupportedError)
+
+        result = cli_runner.invoke(cli.down, ['sky-serve-con*'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.down, ['sky-serve-con*', '-p'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.down, ['-a'], input='n\n')
+        assert isinstance(result.exception, SystemExit), result.exception
+        assert 'Aborted' in result.output
+
+        result = cli_runner.invoke(cli.down, ['-ap'], input='n\n')
+        assert isinstance(result.exception, SystemExit)
+        assert 'Aborted' in result.output
+
+    @pytest.mark.timeout(60)
+    def test_stop_serve_controller(self, _mock_cluster_state,
+                                   _mock_serve_controller):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.stop, [serve.SKY_SERVE_CONTROLLER_NAME])
+        assert result.exit_code == click.UsageError.exit_code
+        assert (
+            f'Stopping controller(s) \'{serve.SKY_SERVE_CONTROLLER_NAME}\' is '
+            'currently not supported' in result.output)
+
+        result = cli_runner.invoke(cli.stop, ['sky-serve-con*'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.stop, ['-a'], input='n\n')
+        assert isinstance(result.exception, SystemExit)
+        assert 'Aborted' in result.output
+
+    @pytest.mark.timeout(60)
+    def test_autostop_serve_controller(self, _mock_cluster_state,
+                                       _mock_serve_controller):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.autostop,
+                                   [serve.SKY_SERVE_CONTROLLER_NAME])
+        assert result.exit_code == click.UsageError.exit_code
+        assert (
+            'Scheduling autostop on controller(s) '
+            f'\'{serve.SKY_SERVE_CONTROLLER_NAME}\' is currently not supported'
+            in result.output)
+
+        result = cli_runner.invoke(cli.autostop, ['sky-serve-con*'])
+        assert not result.exception
+        assert 'Cluster(s) not found' in result.output
+
+        result = cli_runner.invoke(cli.autostop, ['-a'], input='n\n')
+        assert isinstance(result.exception, SystemExit)
+        assert 'Aborted' in result.output
+
+    def test_cancel_on_serve_controller(self, _mock_cluster_state,
+                                        _mock_serve_controller):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.cancel,
+                                   [serve.SKY_SERVE_CONTROLLER_NAME, '-a'])
+        assert result.exit_code == 1
+        assert 'Cancelling the sky serve controller\'s jobs is not allowed.' in str(
+            result.output)
+
+    @pytest.mark.timeout(60)
+    def test_down(self, _mock_db_conn):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.serve_down, ['-a'])
+        assert result.exit_code == 1
+        assert (controller_utils.Controllers.SKY_SERVE_CONTROLLER.value.
+                default_hint_if_non_existent
+                in str(result.output)), (result.exception, result.output,
+                                         result.exc_info)
+
+    @pytest.mark.timeout(60)
+    def test_logs(self, _mock_db_conn):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.serve_logs, ['test', '--controller'])
+        assert controller_utils.Controllers.SKY_SERVE_CONTROLLER.value.default_hint_if_non_existent in str(
+            result.output), (result.exception, result.output, result.exc_info)
+
+    @pytest.mark.timeout(60)
+    def test_status(self, _mock_db_conn):
+        cli_runner = cli_testing.CliRunner()
+        result = cli_runner.invoke(cli.serve_status)
+        assert result.exit_code == 0
+        assert controller_utils.Controllers.SKY_SERVE_CONTROLLER.value.default_hint_if_non_existent in str(
+            result.output), (result.exception, result.output, result.exc_info)


Code review comments:
diff_hunk: @@ -639,7 +640,7 @@ def tail_logs(
                                  'target=CONTROLLER/LOAD_BALANCER.')
     handle = backend_utils.is_controller_accessible(
         controller_type=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
-        stopped_message='No in-progress services.')
+        stopped_message='No live services.')
path: sky/serve/core.py
text: @user1:
```suggestion
   stopped_message=controller_utils.Controllers.SKY_SERVE_CONTROLLER.default_hint_if_non_existent)
```

diff_hunk: @@ -2308,14 +2304,18 @@ def is_controller_accessible(
     error_msg = None
     if controller_status == status_lib.ClusterStatus.STOPPED:
         error_msg = stopped_message
-    elif controller_status is None or handle.head_ip is None:
+    elif controller_status is None or handle is None or handle.head_ip is None:
+        # We check the controller is STOPPED before the check for handle.head_ip
+        # None because when the controller is STOPPED, handle.head_ip can also
+        # be None, but we only want to catch the case when the controller is
+        # being provisioned at the first time and have no head_ip.
         error_msg = non_existent_message
     elif (controller_status == status_lib.ClusterStatus.INIT or
           need_connection_check):
-        # We do not directly check for UP because the controller may be in INIT
-        # state during another `sky spot launch` or `sky serve up`, but still
-        # have head_ip available. In this case, we can still try to ssh into the
-        # controller and fetch the status.
+        # We check the access to controller in INIT state or failed to fetch the
+        # status, as its lock can be hold by another `sky spot launch` or
+        # `sky serve up`, but still have head_ip available. In those cases,
+        # we can allow the access to the controller.
path: sky/backends/backend_utils.py
text: @user1:
```suggestion
        # Check ssh connection if (1) controller is in INIT state, or (2) we failed to fetch the
        # status, both of which can happen when controller's status lock is held by another `sky spot launch` or
        # `sky serve up`. If we have controller's head_ip available and it is ssh-reachable,
        # we can allow access to the controller.
```
Do we mean something like this?

diff_hunk: @@ -796,17 +796,21 @@ def spot_queue(refresh: bool,
             does not exist.
         RuntimeError: if failed to get the spot jobs with ssh.
     """
-
     stop_msg = ''
     if not refresh:
-        stop_msg = 'To view the latest job table: sky spot queue --refresh'
-    controller_status, handle = backend_utils.is_controller_up(
-        controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
-        stopped_message=stop_msg)
+        stop_msg = ('No in-progress spot jobs.\nTo view the latest job table: '
+                    'sky spot queue --refresh')
+    try:
+        handle = backend_utils.is_controller_accessible(
+            controller_type=controller_utils.Controllers.SPOT_CONTROLLER,
+            stopped_message=stop_msg)
path: sky/core.py
text: @user1:
nit: stopped_message to be consistent?