From 763e963774e3c81cd9418178000d7ae69607c8ae Mon Sep 17 00:00:00 2001
From: Romil Bhardwaj <romil.bhardwaj@berkeley.edu>
Date: Thu, 23 Jan 2025 13:03:59 -0800
Subject: [PATCH] [k8s] Show enabled contexts in `sky check` (#4587)

* Show context in check

* rename var
---
 sky/check.py             | 32 +++++++++++++++++++++++++++++++-
 sky/clouds/kubernetes.py |  6 +++---
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/sky/check.py b/sky/check.py
index 1ab92cb1af6..f32e4985079 100644
--- a/sky/check.py
+++ b/sky/check.py
@@ -155,7 +155,8 @@ def get_all_clouds():
         # Pretty print for UX.
         if not quiet:
             enabled_clouds_str = '\n  :heavy_check_mark: '.join(
-                [''] + sorted(all_enabled_clouds))
+                [''] +
+                [_format_enabled_cloud(c) for c in sorted(all_enabled_clouds)])
             rich.print('\n[green]:tada: Enabled clouds :tada:'
                        f'{enabled_clouds_str}[/green]')
 
@@ -222,3 +223,32 @@ def get_cloud_credential_file_mounts(
         r2_credential_mounts = cloudflare.get_credential_file_mounts()
         file_mounts.update(r2_credential_mounts)
     return file_mounts
+
+
+def _format_enabled_cloud(cloud_name: str) -> str:
+    if cloud_name == repr(sky_clouds.Kubernetes()):
+        # Get enabled contexts for Kubernetes
+        existing_contexts = sky_clouds.Kubernetes.existing_allowed_contexts()
+        if not existing_contexts:
+            return cloud_name
+
+        # Check if allowed_contexts is explicitly set in config
+        allowed_contexts = skypilot_config.get_nested(
+            ('kubernetes', 'allowed_contexts'), None)
+
+        # Format the context info with consistent styling
+        if allowed_contexts is not None:
+            contexts_formatted = []
+            for i, context in enumerate(existing_contexts):
+                # TODO: We should use ux_utils.INDENT_SYMBOL and
+                # INDENT_LAST_SYMBOL but, they are formatted for colorama, while
+                # here we are using rich. We should migrate this file to
+                # use colorama as we do in the rest of the codebase.
+                symbol = ('└── ' if i == len(existing_contexts) - 1 else '├── ')
+                contexts_formatted.append(f'\n        {symbol}{context}')
+            context_info = f'Allowed contexts:{"".join(contexts_formatted)}'
+        else:
+            context_info = f'Active context: {existing_contexts[0]}'
+
+        return f'{cloud_name}[/green][dim]\n    └── {context_info}[/dim][green]'
+    return cloud_name
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index f9242bd77aa..3bab7687f48 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -131,7 +131,7 @@ def _log_skipped_contexts_once(cls, skipped_contexts: Tuple[str,
                 'Ignoring these contexts.')
 
     @classmethod
-    def _existing_allowed_contexts(cls) -> List[str]:
+    def existing_allowed_contexts(cls) -> List[str]:
         """Get existing allowed contexts.
 
         If None is returned in the list, it means that we are running in a pod
@@ -175,7 +175,7 @@ def regions_with_offering(cls, instance_type: Optional[str],
                               use_spot: bool, region: Optional[str],
                               zone: Optional[str]) -> List[clouds.Region]:
         del accelerators, zone, use_spot  # unused
-        existing_contexts = cls._existing_allowed_contexts()
+        existing_contexts = cls.existing_allowed_contexts()
 
         regions = []
         for context in existing_contexts:
@@ -591,7 +591,7 @@ def _make(instance_list):
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
         # Test using python API
         try:
-            existing_allowed_contexts = cls._existing_allowed_contexts()
+            existing_allowed_contexts = cls.existing_allowed_contexts()
         except ImportError as e:
             return (False,
                     f'{common_utils.format_exception(e, use_bracket=True)}')