[UX] Add currently activated account for the sky check and check fo…

…r disk-tier options (#2114) * Add --all for the `sky.check` * Better UX * format * change to -v * Add choices for disk tier * Update sky/cli.py Co-authored-by: Tian Xia <[email protected]> --------- Co-authored-by: Tian Xia <[email protected]>
skypilot-org · Jul 10, 2023 · 7262e21 · 7262e21
1 parent c2e9395
commit 7262e21
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 7 deletions.
diff --git a/sky/check.py b/sky/check.py
@@ -9,7 +9,7 @@
 
 
 # TODO(zhwu): add check for a single cloud to improve performance
-def check(quiet: bool = False) -> None:
+def check(quiet: bool = False, verbose: bool = False) -> None:
     echo = (lambda *_args, **_kwargs: None) if quiet else click.echo
     echo('Checking credentials to enable clouds for SkyPilot.')
 
@@ -27,6 +27,10 @@ def check(quiet: bool = False) -> None:
                  ' ' * 10)
         if ok:
             enabled_clouds.append(str(cloud))
+            if verbose:
+                activated_account = cloud.get_current_user_identity_str()
+                if activated_account is not None:
+                    echo(f'    Activated account: {activated_account}')
             if reason is not None:
                 echo(f'    Hint: {reason}')
         else:

diff --git a/sky/cli.py b/sky/cli.py
@@ -224,7 +224,8 @@ def _interactive_node_cli_command(cli_func):
                              help=('OS disk size in GBs.'))
     disk_tier = click.option('--disk-tier',
                              default=None,
-                             type=str,
+                             type=click.Choice(['low', 'medium', 'high'],
+                                               case_sensitive=False),
                              required=False,
                              help=('OS disk tier. Could be one of "low", '
                                    '"medium", "high". Default: medium'))
@@ -1227,7 +1228,7 @@ def cli():
 @click.option(
     '--disk-tier',
     default=None,
-    type=str,
+    type=click.Choice(['low', 'medium', 'high'], case_sensitive=False),
     required=False,
     help=(
         'OS disk tier. Could be one of "low", "medium", "high". Default: medium'
@@ -3004,8 +3005,13 @@ def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 
 
 @cli.command()
+@click.option('--verbose',
+              '-v',
+              is_flag=True,
+              default=False,
+              help='Show the activated account for each cloud.')
 @usage_lib.entrypoint
-def check():
+def check(verbose: bool):
     """Check which clouds are available to use.
 
     This checks access credentials for all clouds supported by SkyPilot. If a
@@ -3015,7 +3021,7 @@ def check():
     The enabled clouds are cached and form the "search space" to be considered
     for each task.
     """
-    sky_check.check()
+    sky_check.check(verbose=verbose)
 
 
 @cli.command()
@@ -3399,7 +3405,7 @@ def spot():
 @click.option(
     '--disk-tier',
     default=None,
-    type=str,
+    type=click.Choice(['low', 'medium', 'high'], case_sensitive=False),
     required=False,
     help=(
         'OS disk tier. Could be one of "low", "medium", "high". Default: medium'
@@ -3859,7 +3865,7 @@ def bench():
 @click.option(
     '--disk-tier',
     default=None,
-    type=str,
+    type=click.Choice(['low', 'medium', 'high'], case_sensitive=False),
     required=False,
     help=(
         'OS disk tier. Could be one of "low", "medium", "high". Default: medium'

diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
@@ -622,6 +622,14 @@ def get_current_user_identity(cls) -> Optional[List[str]]:
                 ) from None
         return user_ids
 
+    @classmethod
+    def get_current_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_current_user_identity()
+        if user_identity is None:
+            return None
+        identity_str = f'{user_identity[0]} [account={user_identity[1]}]'
+        return identity_str
+
     def get_credential_file_mounts(self) -> Dict[str, str]:
         # The credentials file should not be uploaded if the user identity is
         # not SHARED_CREDENTIALS_FILE, since we cannot be sure if the currently

diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
@@ -444,6 +444,13 @@ def get_current_user_identity(cls) -> Optional[List[str]]:
                     'Failed to get Azure project ID.') from e
         return [f'{account_email} [subscription_id={project_id}]']
 
+    @classmethod
+    def get_current_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_current_user_identity()
+        if user_identity is None:
+            return None
+        return user_identity[0]
+
     @classmethod
     def get_project_id(cls, dryrun: bool = False) -> str:
         if dryrun:

diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
@@ -354,6 +354,14 @@ def get_current_user_identity(cls) -> Optional[List[str]]:
         """
         return None
 
+    @classmethod
+    def get_current_user_identity_str(cls) -> Optional[str]:
+        """Returns a user friendly representation of the current identity."""
+        user_identity = cls.get_current_user_identity()
+        if user_identity is None:
+            return None
+        return ', '.join(user_identity)
+
     def get_credential_file_mounts(self) -> Dict[str, str]:
         """Returns the files necessary to access this cloud.
 

diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
@@ -720,6 +720,13 @@ def get_current_user_identity(cls) -> Optional[List[str]]:
                 ) from e
         return [f'{account} [project_id={project_id}]']
 
+    @classmethod
+    def get_current_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_current_user_identity()
+        if user_identity is None:
+            return None
+        return user_identity[0].replace('\n', '')
+
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, 'gcp')