Add user identity to cluster status to avoid leakage when switching account

sky/backends/backend_utils.py

@@ -14,7 +14,7 @@
 import threading
 import time
 import typing
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Optional, Set, Tuple, Type
 import uuid
 
 import colorama
@@ -1644,7 +1644,9 @@ def check_owner_identity(cluster_name: str) -> None:
 
     Raises:
         exceptions.ClusterOwnerIdentityMismatchError: if the current user is not the
-            same as the user who created the cluster.
+          same as the user who created the cluster.
+        sky.exceptions.CloudUserIdentityError: if failed to get the activated user
+          identity.
     """
     if env_options.Options.SKIP_CLOUD_IDENTITY_CHECK.get():
         return
@@ -1664,18 +1666,19 @@ def check_owner_identity(cluster_name: str) -> None:
     # The user identity can be None, if the cluster is created by an older
     # version of SkyPilot. In that case, we set the user identity to the
     # current one.
-    # NOTE: a user upgrades SkyPilot and switches to a new cloud identity
-    # immediately without `sky status --refresh` first, can cause a leakage
-    # of the existing cluster.
+    # NOTE: a user who upgrades SkyPilot and switches to a new cloud identity
+    # immediately without `sky status --refresh` first, will cause a leakage
+    # of the existing cluster. We deem this an acceptable tradeoff mainly
+    # because multi-identity is not common (at least at the moment).
     if owner_identity is None:
         global_user_state.set_owner_identity_for_cluster(
             cluster_name, current_user_identity)
     elif owner_identity != current_user_identity:
         with ux_utils.print_exception_no_traceback():
             raise exceptions.ClusterOwnerIdentityMismatchError(
-                f'The cluster {cluster_name!r} (on {cloud}) is created by the account '
-                f'{owner_identity!r}, but the activated one is {current_user_identity!r}.'
-            )
+                f'Cluster {cluster_name!r} ({cloud}) is owned by account '
+                f'{owner_identity!r}, but the currently activated account '
+                f'is {current_user_identity!r}.')
 
 
 def _get_cluster_status_via_cloud_cli(
@@ -1800,9 +1803,11 @@ def _update_cluster_identity_and_status(
 
     Raises:
         sky.exceptions.ClusterOwnerIdentityMismatchError: the cluster's owner
-        identity mismatches the current cloud user.
+          identity mismatches the current cloud user.
+        sky.exceptions.CloudUserIdentityError: if failed to get the activated user
+          identity.
         sky.exceptions.ClusterStatusFetchingError: the cluster status cannot be
-        fetched from the cloud provider.
+          fetched from the cloud provider.
     """
     check_owner_identity(cluster_name)
 
@@ -1829,8 +1834,19 @@ def refresh_cluster_status_handle(
     *,
     force_refresh: bool = False,
     acquire_per_cluster_status_lock: bool = True,
+    suppress_error: bool = False
 ) -> Tuple[Optional[global_user_state.ClusterStatus],
            Optional[backends.Backend.ResourceHandle]]:
+    """Refresh the cluster status and return the status and handle.
+
+    Raises:
+        sky.exceptions.ClusterOwnerIdentityMismatchError: the cluster's owner
+          identity mismatches the current cloud user.
+        sky.exceptions.CloudUserIdentityError: if failed to get the activated user
+          identity.
+        sky.exceptions.ClusterStatusFetchingError: the cluster status cannot be
+          fetched from the cloud provider.
+    """
     record = global_user_state.get_cluster_from_name(cluster_name)
     if record is None:
         return None, None
@@ -1839,14 +1855,87 @@ def refresh_cluster_status_handle(
     if isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
         use_spot = handle.launched_resources.use_spot
         if (force_refresh or record['autostop'] >= 0 or use_spot):
-            record = _update_cluster_identity_and_status(
-                cluster_name,
-                acquire_per_cluster_status_lock=acquire_per_cluster_status_lock)
-            if record is None:
-                return None, None
+            try:
+                record = _update_cluster_identity_and_status(
+                    cluster_name,
+                    acquire_per_cluster_status_lock=
+                    acquire_per_cluster_status_lock)
+                if record is None:
+                    return None, None
+            except (exceptions.ClusterOwnerIdentityMismatchError,
+                    exceptions.CloudUserIdentityError,
+                    exceptions.ClusterStatusFetchingError) as e:
+                if suppress_error:
+                    logger.debug(
+                        f'Failed to refresh cluster {cluster_name!r} due to {e}'
+                    )
+                    return None, None
+                raise
     return record['status'], record['handle']
 
 
+def check_cluster_available(
+    cluster_name: str,
+    operation: str,
+    *,
+    expected_backend: Optional[Type[backends.Backend]] = None
+) -> backends.Backend.ResourceHandle:
+    """Check if the cluster is available.
+
+    Raises:
+        exceptions.ClusterNotUpError: if the cluster is not UP.
+        exceptions.NotSupportedError: if the cluster is not based on
+          CloudVmRayBackend
+    """
+    try:
+        cluster_status, handle = refresh_cluster_status_handle(cluster_name)
+    except (exceptions.ClusterOwnerIdentityMismatchError,
+            exceptions.CloudUserIdentityError,
+            exceptions.ClusterStatusFetchingError) as e:
+        # Failed to refresh the cluster status is not fatal error as the callers
+        # can still be done by only using ssh, but the ssh can hang if the
+        # cluster is not up (e.g., autostopped).
+        ux_utils.console_newline()
+        logger.warning(
+            f'Failed to refresh the cluster status, it is not fatal, but '
+            f'{operation} cluster {cluster_name!r} might hang if the cluster '
+            'is not up.\n'
+            f'Detailed reason: {e}')
+        record = global_user_state.get_cluster_from_name(cluster_name)
+        cluster_status, handle = record['status'], record['handle']
+
+    if handle is None:
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.ClusterNotUpError(
+                f'{colorama.Fore.YELLOW}Cluster {cluster_name!r} does not '
+                f'exist... skipped.{colorama.Style.RESET_ALL}')
+    backend = get_backend_from_handle(handle)
+    if expected_backend is not None and not isinstance(backend,
+                                                       expected_backend):
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.NotSupportedError(
+                f'{colorama.Fore.YELLOW}{operation} cluster '
+                f'{cluster_name!r} (status: {cluster_status.value})... skipped'
+                f'{colorama.Style.RESET_ALL}')
+    if cluster_status != global_user_state.ClusterStatus.UP:
+        if onprem_utils.check_if_local_cloud(cluster_name):
+            raise exceptions.ClusterNotUpError(
+                constants.UNINITIALIZED_ONPREM_CLUSTER_MESSAGE.format(
+                    cluster_name))
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.ClusterNotUpError(
+                f'{colorama.Fore.YELLOW}{operation} cluster {cluster_name!r} '
+                f'(status: {cluster_status.value})... skipped.'
+                f'{colorama.Style.RESET_ALL}')
+
+    if handle.head_ip is None:
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.ClusterNotUpError(
+                f'Cluster {cluster_name!r} has been stopped or not properly '
+                'set up. Please re-launch it with `sky start`.')
+    return handle
+
+
 class CloudFilter(enum.Enum):
     # Filter for all types of clouds.
     ALL = 'all'

sky/cli.py

@@ -3052,7 +3052,7 @@ def spot_cancel(name: Optional[str], job_ids: Tuple[int], all: bool, yes: bool):
 
     """
 
-    _, handle = _is_spot_controller_up(
+    _, handle = spot_lib.is_spot_controller_up(
         'All managed spot jobs should have finished.')
     if handle is None:
         return

sky/core.py

@@ -13,7 +13,6 @@
 from sky import sky_logging
 from sky import spot
 from sky.backends import backend_utils
-from sky.backends import onprem_utils
 from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.usage import usage_lib
@@ -356,10 +355,23 @@ def autostop(
         raise exceptions.NotSupportedError(
             f'{operation} sky reserved cluster {cluster_name!r} '
             f'is not supported.')
-    (cluster_status,
-     handle) = backend_utils.refresh_cluster_status_handle(cluster_name)
-    if handle is None:
-        raise ValueError(f'Cluster {cluster_name!r} does not exist.')
+    try:
+        handle = backend_utils.check_cluster_available(
+            cluster_name,
+            operation,
+            expected_backend=backends.CloudVmRayBackend)
+    except exceptions.ClusterNotUpError as e:
+        with ux_utils.print_exception_no_traceback():
+            e.message += (
+                f'\n  auto{option_str} can only be set/unset for '
+                f'{global_user_state.ClusterStatus.UP.value} clusters.')
+            raise e from None
+    except exceptions.NotSupportedError as e:
+        with ux_utils.print_exception_no_traceback():
+            e.message += (f'\n  auto{option_str} is only supported by backend: '
+                          f'{backends.CloudVmRayBackend.NAME}')
+            raise e from None
+
     if tpu_utils.is_tpu_vm_pod(handle.launched_resources):
         # Reference:
         # https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm#stopping_a_with_gcloud  # pylint: disable=line-too-long
@@ -368,21 +380,6 @@ def autostop(
             'is not supported.')
 
     backend = backend_utils.get_backend_from_handle(handle)
-    if not isinstance(backend, backends.CloudVmRayBackend):
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.NotSupportedError(
-                f'{colorama.Fore.YELLOW}{operation} cluster '
-                f'{cluster_name!r}... skipped{colorama.Style.RESET_ALL}'
-                f'\n  auto{option_str} is only supported by backend: '
-                f'{backends.CloudVmRayBackend.NAME}')
-    if cluster_status != global_user_state.ClusterStatus.UP:
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ClusterNotUpError(
-                f'{colorama.Fore.YELLOW}{operation} cluster '
-                f'{cluster_name!r} (status: {cluster_status.value})... skipped'
-                f'{colorama.Style.RESET_ALL}'
-                f'\n  auto{option_str} can only be set/unset for '
-                f'{global_user_state.ClusterStatus.UP.value} clusters.')
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
     backend.set_autostop(handle, idle_minutes, down)
 
@@ -392,39 +389,6 @@ def autostop(
 # ==================
 
 
-def _check_cluster_available(cluster_name: str,
-                             operation: str) -> backends.Backend.ResourceHandle:
-    """Check if the cluster is available."""
-    cluster_status, handle = backend_utils.refresh_cluster_status_handle(
-        cluster_name)
-    if handle is None:
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ClusterNotUpError(
-                f'{colorama.Fore.YELLOW}Cluster {cluster_name!r} does not '
-                f'exist; skipped.{colorama.Style.RESET_ALL}')
-    backend = backend_utils.get_backend_from_handle(handle)
-    if isinstance(backend, backends.LocalDockerBackend):
-        # LocalDockerBackend does not support job queues
-        raise exceptions.NotSupportedError(
-            f'Cluster {cluster_name} with LocalDockerBackend does '
-            f'not support {operation}.')
-    if cluster_status != global_user_state.ClusterStatus.UP:
-        if onprem_utils.check_if_local_cloud(cluster_name):
-            raise exceptions.ClusterNotUpError(
-                constants.UNINITIALIZED_ONPREM_CLUSTER_MESSAGE.format(
-                    cluster_name))
-        raise exceptions.ClusterNotUpError(
-            f'{colorama.Fore.YELLOW}Cluster {cluster_name!r} is not up '
-            f'(status: {cluster_status.value}); skipped.'
-            f'{colorama.Style.RESET_ALL}')
-
-    if handle.head_ip is None:
-        raise exceptions.ClusterNotUpError(
-            f'Cluster {cluster_name!r} has been stopped or not properly set up.'
-            ' Please re-launch it with `sky start`.')
-    return handle
-
-
 @usage_lib.entrypoint
 def queue(cluster_name: str,
           skip_finished: bool = False,
@@ -460,7 +424,10 @@ def queue(cluster_name: str,
         username = None
     code = job_lib.JobLibCodeGen.get_job_queue(username, all_jobs)
 
-    handle = _check_cluster_available(cluster_name, 'getting the job queue')
+    handle = backend_utils.check_cluster_available(
+        cluster_name,
+        'getting the job queue',
+        expected_backend=backends.CloudVmRayBackend)
     backend = backend_utils.get_backend_from_handle(handle)
 
     returncode, jobs_payload, stderr = backend.run_on_head(handle,
@@ -500,7 +467,10 @@ def cancel(cluster_name: str,
         cluster_name, operation_str='Cancelling jobs')
 
     # Check the status of the cluster.
-    handle = _check_cluster_available(cluster_name, 'cancelling jobs')
+    handle = backend_utils.check_cluster_available(
+        cluster_name,
+        'cancelling jobs',
+        expected_backend=backends.CloudVmRayBackend)
     backend = backend_utils.get_backend_from_handle(handle)
 
     if all:
@@ -532,7 +502,10 @@ def tail_logs(cluster_name: str,
         sky.exceptions.NotSupportedError: the feature is not supported.
     """
     # Check the status of the cluster.
-    handle = _check_cluster_available(cluster_name, 'tailing logs')
+    handle = backend_utils.check_cluster_available(
+        cluster_name,
+        'tailing logs',
+        expected_backend=backends.CloudVmRayBackend)
     backend = backend_utils.get_backend_from_handle(handle)
 
     job_str = f'job {job_id}'
@@ -561,7 +534,10 @@ def download_logs(
         Dict[str, str]: a mapping of job_id to local log path.
     """
     # Check the status of the cluster.
-    handle = _check_cluster_available(cluster_name, 'downloading logs')
+    handle = backend_utils.check_cluster_available(
+        cluster_name,
+        'downloading logs',
+        expected_backend=backends.CloudVmRayBackend)
     backend = backend_utils.get_backend_from_handle(handle)
 
     if job_ids is not None and len(job_ids) == 0:
@@ -593,7 +569,10 @@ def job_status(
         {None: None}.
     """
     # Check the status of the cluster.
-    handle = _check_cluster_available(cluster_name, 'getting job status')
+    handle = backend_utils.check_cluster_available(
+        cluster_name,
+        'getting job status',
+        expected_backend=backends.CloudVmRayBackend)
     backend = backend_utils.get_backend_from_handle(handle)
 
     if job_ids is not None and len(job_ids) == 0:
@@ -613,26 +592,6 @@ def job_status(
 # =======================
 
 
-def _is_spot_controller_up(
-    stopped_message: str,
-) -> Tuple[Optional[global_user_state.ClusterStatus],
-           Optional[backends.Backend.ResourceHandle]]:
-    controller_status, handle = backend_utils.refresh_cluster_status_handle(
-        spot.SPOT_CONTROLLER_NAME, force_refresh=True)
-    if controller_status is None:
-        print('No managed spot job has been run.')
-    elif controller_status != global_user_state.ClusterStatus.UP:
-        msg = (f'Spot controller {spot.SPOT_CONTROLLER_NAME} '
-               f'is {controller_status.value}.')
-        if controller_status == global_user_state.ClusterStatus.STOPPED:
-            msg += f'\n{stopped_message}'
-        if controller_status == global_user_state.ClusterStatus.INIT:
-            msg += '\nPlease wait for the controller to be ready.'
-        print(msg)
-        handle = None
-    return controller_status, handle
-
-
 @usage_lib.entrypoint
 def spot_status(refresh: bool) -> List[Dict[str, Any]]:
     """[Deprecated] (alias of spot_queue) Get statuses of managed spot jobs."""
@@ -672,7 +631,7 @@ def spot_queue(refresh: bool) -> List[Dict[str, Any]]:
     stop_msg = ''
     if not refresh:
         stop_msg = 'To view the latest job table: sky spot queue --refresh'
-    controller_status, handle = _is_spot_controller_up(stop_msg)
+    controller_status, handle = spot.is_spot_controller_up(stop_msg)
 
     if controller_status is None:
         return []
@@ -725,10 +684,12 @@ def spot_cancel(name: Optional[str] = None,
         RuntimeError: failed to cancel the job.
     """
     job_ids = [] if job_ids is None else job_ids
-    _, handle = _is_spot_controller_up(
+    _, handle = spot.is_spot_controller_up(
         'All managed spot jobs should have finished.')
     if handle is None or handle.head_ip is None:
-        raise exceptions.ClusterNotUpError('All jobs finished.')
+        # The error message is already printed in spot.is_spot_controller_up.
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.ClusterNotUpError()
 
     job_id_str = ','.join(map(str, job_ids))
     if sum([len(job_ids) > 0, name is not None, all]) != 1:
@@ -778,7 +739,7 @@ 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 = _is_spot_controller_up(
+    controller_status, handle = spot.is_spot_controller_up(
         'Please restart the spot controller with '
         f'`sky start {spot.SPOT_CONTROLLER_NAME}`.')
     if handle is None or handle.head_ip is None: