Repository: skypilot-org/skypilot
Pull request #5255: Fix querying kubernetes nodes at client-side
Author: login: aylei
Created: 2025-04-17T04:33:45Z | Merged: 2025-04-17T06:20:52Z

<!-- Describe the changes in this PR -->

close https://github.com/skypilot-org/skypilot/issues/5245

**NOTE**: this is a break change so the API version is bumped. Requires an upgrade for both client-server. I've actually implemented a compatible version in https://github.com/skypilot-org/skypilot/pull/5255/commits/3e03113bf7522e1cb8a5a3032dfab8f5fc500315. But considering the maintenance burden, I think it is a good time to do some break changes before we release 0.9.0 so I remove the code about compatibility. How do you like? @cg505 @romilbhardwaj @Michaelvll 

Behavior with mocked labels:

```
sky show-gpus
Kubernetes GPUs
GPU   REQUESTABLE_QTY_PER_NODE  TOTAL_GPUS  TOTAL_FREE_GPUS
H100  1, 2, 4, 8                8           8
L4    1, 2, 4, 8                8           8

Kubernetes per node accelerator availability (Note: Multi-host TPUs are detected and excluded from the display as multi-host TPUs are not supported.)
NODE_NAME                                     GPU_NAME  TOTAL_GPUS  FREE_GPUS
gke-skypilotalpha-default-pool-ff931856-y51p  None      0           0
gke-skypilotalpha-largecpu-05dae726-0f7l      H100      8           8
gke-skypilotalpha-largecpu-05dae726-hyv6      L4        8           8
```

<!-- 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: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->



Diff:
diff --git a/sky/cli.py b/sky/cli.py
index fb93694a545..965b29488e4 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -3463,15 +3463,14 @@ def _get_kubernetes_realtime_gpu_table(
             ])
         return realtime_gpu_table
 
-    # TODO(zhwu): this needs to run on remote server.
-    def _get_kubernetes_node_info_table(context: Optional[str]):
+    def _format_kubernetes_node_info(context: Optional[str]):
         node_table = log_utils.create_table(
             ['NODE_NAME', 'GPU_NAME', 'TOTAL_GPUS', 'FREE_GPUS'])
 
-        no_permissions_str = '<no permissions>'
-        node_info_dict = sdk.stream_and_get(
+        nodes_info = sdk.stream_and_get(
             sdk.kubernetes_node_info(context=context))
-        for node_name, node_info in node_info_dict.items():
+        no_permissions_str = '<no permissions>'
+        for node_name, node_info in nodes_info.node_info_dict.items():
             available = node_info.free[
                 'accelerators_available'] if node_info.free[
                     'accelerators_available'] != -1 else no_permissions_str
@@ -3479,7 +3478,14 @@ def _get_kubernetes_node_info_table(context: Optional[str]):
                 node_name, node_info.accelerator_type,
                 node_info.total['accelerator_count'], available
             ])
-        return node_table
+        k8s_per_node_acc_message = (
+            'Kubernetes per node accelerator availability ')
+        if nodes_info.hint:
+            k8s_per_node_acc_message += nodes_info.hint
+        return (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+                f'{k8s_per_node_acc_message}'
+                f'{colorama.Style.RESET_ALL}\n'
+                f'{node_table.get_string()}')
 
     def _output() -> Generator[str, None, None]:
         gpu_table = log_utils.create_table(
@@ -3527,22 +3533,8 @@ def _output() -> Generator[str, None, None]:
                            f'Kubernetes GPUs {context_str}'
                            f'{colorama.Style.RESET_ALL}\n')
                     yield from k8s_realtime_table.get_string()
-                    k8s_node_table = _get_kubernetes_node_info_table(context)
                     yield '\n\n'
-                    # TODO(Doyoung): Update the message with the multi-host TPU
-                    # support.
-                    k8s_per_node_acc_message = (
-                        'Kubernetes per node accelerator availability ')
-                    if kubernetes_utils.multi_host_tpu_exists_in_cluster(
-                            context):
-                        k8s_per_node_acc_message += (
-                            '(Note: Multi-host TPUs are detected and excluded '
-                            'from the display as multi-host TPUs are not '
-                            'supported.)')
-                    yield (f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
-                           f'{k8s_per_node_acc_message}'
-                           f'{colorama.Style.RESET_ALL}\n')
-                    yield from k8s_node_table.get_string()
+                    yield _format_kubernetes_node_info(context)
                 if kubernetes_autoscaling:
                     k8s_messages += (
                         '\n' + kubernetes_utils.KUBERNETES_AUTOSCALER_NOTE)
diff --git a/sky/client/sdk.py b/sky/client/sdk.py
index bef9c244205..9c98d285dce 100644
--- a/sky/client/sdk.py
+++ b/sky/client/sdk.py
@@ -1408,8 +1408,8 @@ def kubernetes_node_info(
         The request ID of the Kubernetes node info request.
 
     Request Returns:
-        Dict[str, KubernetesNodeInfo]: Dictionary containing the node name as
-            key and the KubernetesNodeInfo object as value
+        KubernetesNodesInfo: A model that contains the node info map and other
+            information.
     """
     body = payloads.KubernetesNodeInfoRequestBody(context=context)
     response = requests.post(
diff --git a/sky/models.py b/sky/models.py
index 898725419a1..a5fb28e9a6e 100644
--- a/sky/models.py
+++ b/sky/models.py
@@ -28,3 +28,31 @@ class KubernetesNodeInfo:
     # Resources available on the node. E.g., {'nvidia.com/gpu': '2'}
     total: Dict[str, int]
     free: Dict[str, int]
+
+
+@dataclasses.dataclass
+class KubernetesNodesInfo:
+    """Dataclass to store Kubernetes node info map."""
+    # The nodes in the cluster, keyed by node name.
+    node_info_dict: Dict[str, KubernetesNodeInfo]
+    # Additional hint for the node info.
+    hint: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            'node_info_dict': {
+                node_name: dataclasses.asdict(node_info)
+                for node_name, node_info in self.node_info_dict.items()
+            },
+            'hint': self.hint,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'KubernetesNodesInfo':
+        return cls(
+            node_info_dict={
+                node_name: KubernetesNodeInfo(**node_info)
+                for node_name, node_info in data['node_info_dict'].items()
+            },
+            hint=data['hint'],
+        )
diff --git a/sky/provision/kubernetes/utils.py b/sky/provision/kubernetes/utils.py
index d84fdb128b0..4e78b24c8ee 100644
--- a/sky/provision/kubernetes/utils.py
+++ b/sky/provision/kubernetes/utils.py
@@ -2543,9 +2543,15 @@ def get_unlabeled_accelerator_nodes(context: Optional[str] = None) -> List[Any]:
 
 
 def get_kubernetes_node_info(
-        context: Optional[str] = None) -> Dict[str, models.KubernetesNodeInfo]:
+        context: Optional[str] = None) -> models.KubernetesNodesInfo:
     """Gets the resource information for all the nodes in the cluster.
 
+    This function returns a model with node info map as a nested field. This
+    allows future extensions while keeping the client-server compatibility,
+    e.g. when adding a new field to the model, the legacy clients will not be
+    affected and new clients can opt-in new behavior if the new field is
+    presented.
+
     Currently only GPU resources are supported. The function returns the total
     number of GPUs available on the node and the number of free GPUs on the
     node.
@@ -2554,8 +2560,8 @@ def get_kubernetes_node_info(
     namespaces, the function will return free GPUs as -1.
 
     Returns:
-        Dict[str, KubernetesNodeInfo]: Dictionary containing the node name as
-            key and the KubernetesNodeInfo object as value
+        KubernetesNodesInfo: A model that contains the node info map and other
+            information.
     """
     nodes = get_kubernetes_nodes(context=context)
     # Get the pods to get the real-time resource usage
@@ -2574,6 +2580,7 @@ def get_kubernetes_node_info(
         label_keys = lf.get_label_keys()
 
     node_info_dict: Dict[str, models.KubernetesNodeInfo] = {}
+    has_multi_host_tpu = False
 
     for node in nodes:
         accelerator_name = None
@@ -2610,6 +2617,7 @@ def get_kubernetes_node_info(
         # TODO(Doyoung): Remove the logic when adding support for
         # multi-host TPUs.
         if is_multi_host_tpu(node.metadata.labels):
+            has_multi_host_tpu = True
             continue
 
         node_info_dict[node.metadata.name] = models.KubernetesNodeInfo(
@@ -2617,8 +2625,15 @@ def get_kubernetes_node_info(
             accelerator_type=accelerator_name,
             total={'accelerator_count': int(accelerator_count)},
             free={'accelerators_available': int(accelerators_available)})
+    hint = ''
+    if has_multi_host_tpu:
+        hint = ('(Note: Multi-host TPUs are detected and excluded from the '
+                'display as multi-host TPUs are not supported.)')
 
-    return node_info_dict
+    return models.KubernetesNodesInfo(
+        node_info_dict=node_info_dict,
+        hint=hint,
+    )
 
 
 def to_label_selector(tags):
@@ -2865,15 +2880,6 @@ def is_multi_host_tpu(node_metadata_labels: dict) -> bool:
     return False
 
 
-def multi_host_tpu_exists_in_cluster(context: Optional[str] = None) -> bool:
-    """Checks if there exists a multi-host TPU within the cluster."""
-    nodes = get_kubernetes_nodes(context=context)
-    for node in nodes:
-        if is_multi_host_tpu(node.metadata.labels):
-            return True
-    return False
-
-
 @dataclasses.dataclass
 class KubernetesSkyPilotClusterInfo:
     cluster_name_on_cloud: str
diff --git a/sky/server/constants.py b/sky/server/constants.py
index 8f328400ab5..e3e3ae332e6 100644
--- a/sky/server/constants.py
+++ b/sky/server/constants.py
@@ -5,7 +5,7 @@
 # API server version, whenever there is a change in API server that requires a
 # restart of the local API server or error out when the client does not match
 # the server version.
-API_VERSION = '3'
+API_VERSION = '4'
 
 # Prefix for API request names.
 REQUEST_NAME_PREFIX = 'sky.'
diff --git a/sky/server/requests/serializers/decoders.py b/sky/server/requests/serializers/decoders.py
index b0d2eb9d862..8f6de574658 100644
--- a/sky/server/requests/serializers/decoders.py
+++ b/sky/server/requests/serializers/decoders.py
@@ -188,8 +188,5 @@ def decode_job_status(
 
 @register_decoders('kubernetes_node_info')
 def decode_kubernetes_node_info(
-        return_value: Dict[str, Any]) -> Dict[str, models.KubernetesNodeInfo]:
-    return {
-        node_name: models.KubernetesNodeInfo(**node_info)
-        for node_name, node_info in return_value.items()
-    }
+        return_value: Dict[str, Any]) -> models.KubernetesNodesInfo:
+    return models.KubernetesNodesInfo.from_dict(return_value)
diff --git a/sky/server/requests/serializers/encoders.py b/sky/server/requests/serializers/encoders.py
index 0683d6ecdff..8fb05e0a1fc 100644
--- a/sky/server/requests/serializers/encoders.py
+++ b/sky/server/requests/serializers/encoders.py
@@ -159,8 +159,5 @@ def encode_job_status(return_value: Dict[int, Any]) -> Dict[int, str]:
 
 @register_encoder('kubernetes_node_info')
 def encode_kubernetes_node_info(
-        return_value: Dict[str, 'models.KubernetesNodeInfo']) -> Dict[str, Any]:
-    return {
-        node_name: dataclasses.asdict(node_info)
-        for node_name, node_info in return_value.items()
-    }
+        return_value: 'models.KubernetesNodesInfo') -> Dict[str, Any]:
+    return return_value.to_dict()
diff --git a/sky/server/server.py b/sky/server/server.py
index 39435b1eeb8..c219acdb8e5 100644
--- a/sky/server/server.py
+++ b/sky/server/server.py
@@ -210,7 +210,7 @@ async def kubernetes_node_info(
         request: fastapi.Request,
         kubernetes_node_info_body: payloads.KubernetesNodeInfoRequestBody
 ) -> None:
-    """Gets Kubernetes node information."""
+    """Gets Kubernetes nodes information and hints."""
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='kubernetes_node_info',


Reviews:
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:47:43Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:48:46Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:51:09Z
user: login: Michaelvll
body: Thanks @aylei! LGTM.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-17T05:01:16Z
user: login: Michaelvll
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T05:04:48Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T05:19:26Z
user: login: zpoint
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T06:22:19Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T06:30:47Z

Review comments:
diff_hunk: @@ -3463,23 +3463,29 @@ def _get_kubernetes_realtime_gpu_table(
             ])
         return realtime_gpu_table
 
-    # TODO(zhwu): this needs to run on remote server.
path: sky/cli.py
user: login: aylei
body: I removed this TODO since this function only does formatting, which varies in different clients, e.g. CLI and Web UI. Would be good to keep this at client-side. Wdyt? @Michaelvll 

cc @DanielZhangQD 
created_at: 2025-04-17T04:46:31Z
author_association: COLLABORATOR
diff_hunk: @@ -5,7 +5,7 @@
 # API server version, whenever there is a change in API server that requires a
 # restart of the local API server or error out when the client does not match
 # the server version.
-API_VERSION = '3'
+API_VERSION = '4'
path: sky/server/constants.py
user: login: aylei
body: Break change: the response struct of kubernetes_node_info was changed
created_at: 2025-04-17T04:48:43Z
author_association: COLLABORATOR
diff_hunk: @@ -3463,23 +3463,29 @@ def _get_kubernetes_realtime_gpu_table(
             ])
         return realtime_gpu_table
 
-    # TODO(zhwu): this needs to run on remote server.
path: sky/cli.py
user: login: aylei
body: Also renamed the function to `_format_kubernetes_node_info` for clarity
created_at: 2025-04-17T04:51:09Z
author_association: COLLABORATOR
diff_hunk: @@ -3463,23 +3463,29 @@ def _get_kubernetes_realtime_gpu_table(
             ])
         return realtime_gpu_table
 
-    # TODO(zhwu): this needs to run on remote server.
path: sky/cli.py
user: login: Michaelvll
body: Ahh, this TODO was mainly to reduce the RTT by node_info_dict = sdk.stream_and_get(sdk.kubernetes_node_info(context=context))` below. It is fine to remove : )
created_at: 2025-04-17T05:00:12Z
author_association: COLLABORATOR
diff_hunk: @@ -3527,22 +3533,8 @@ def _output() -> Generator[str, None, None]:
                            f'Kubernetes GPUs {context_str}'
                            f'{colorama.Style.RESET_ALL}\n')
                     yield from k8s_realtime_table.get_string()
-                    k8s_node_table = _get_kubernetes_node_info_table(context)
                     yield '\n\n'
-                    # TODO(Doyoung): Update the message with the multi-host TPU
-                    # support.
-                    k8s_per_node_acc_message = (
-                        'Kubernetes per node accelerator availability ')
-                    if kubernetes_utils.multi_host_tpu_exists_in_cluster(
path: sky/cli.py
user: login: Michaelvll
body: Can we add a test to avoid this kind of regression in the future?
created_at: 2025-04-17T05:04:43Z
author_association: COLLABORATOR
diff_hunk: @@ -3527,22 +3533,8 @@ def _output() -> Generator[str, None, None]:
                            f'Kubernetes GPUs {context_str}'
                            f'{colorama.Style.RESET_ALL}\n')
                     yield from k8s_realtime_table.get_string()
-                    k8s_node_table = _get_kubernetes_node_info_table(context)
                     yield '\n\n'
-                    # TODO(Doyoung): Update the message with the multi-host TPU
-                    # support.
-                    k8s_per_node_acc_message = (
-                        'Kubernetes per node accelerator availability ')
-                    if kubernetes_utils.multi_host_tpu_exists_in_cluster(
path: sky/cli.py
user: login: aylei
body: Maybe get rid of all cloud permissions at client-side in smoke test when running with `--remote-server`, cc @zpoint wdyt?
created_at: 2025-04-17T05:19:26Z
author_association: COLLABORATOR
diff_hunk: @@ -3527,22 +3533,8 @@ def _output() -> Generator[str, None, None]:
                            f'Kubernetes GPUs {context_str}'
                            f'{colorama.Style.RESET_ALL}\n')
                     yield from k8s_realtime_table.get_string()
-                    k8s_node_table = _get_kubernetes_node_info_table(context)
                     yield '\n\n'
-                    # TODO(Doyoung): Update the message with the multi-host TPU
-                    # support.
-                    k8s_per_node_acc_message = (
-                        'Kubernetes per node accelerator availability ')
-                    if kubernetes_utils.multi_host_tpu_exists_in_cluster(
path: sky/cli.py
user: login: zpoint
body: Perhaps it's better if we can create test credentials directly in the test case instead of relying on the real(buildkite) system environment.
created_at: 2025-04-17T06:22:19Z
author_association: COLLABORATOR
diff_hunk: @@ -3527,22 +3533,8 @@ def _output() -> Generator[str, None, None]:
                            f'Kubernetes GPUs {context_str}'
                            f'{colorama.Style.RESET_ALL}\n')
                     yield from k8s_realtime_table.get_string()
-                    k8s_node_table = _get_kubernetes_node_info_table(context)
                     yield '\n\n'
-                    # TODO(Doyoung): Update the message with the multi-host TPU
-                    # support.
-                    k8s_per_node_acc_message = (
-                        'Kubernetes per node accelerator availability ')
-                    if kubernetes_utils.multi_host_tpu_exists_in_cluster(
path: sky/cli.py
user: login: aylei
body: Yes a specific test case will be helpful and I think we also want to cover https://github.com/skypilot-org/skypilot/issues/5259 in all cases we have currently, to avoid the client accessing any cloud specific code path
created_at: 2025-04-17T06:30:47Z
author_association: COLLABORATOR

PR comments:
user: login: aylei
created_at: 2025-04-17T05:21:45Z
body: Merge to unblock 0.9.0 first, unit test will be added after addressing https://github.com/skypilot-org/skypilot/issues/5256
author_association: COLLABORATOR
user: login: aylei
created_at: 2025-04-17T05:22:30Z
body: /quicktest-core

The break change in compatibility test is expected: https://buildkite.com/skypilot-1/quicktest-core/builds/282#0196423d-94af-4954-b514-2bfb3d4f925f

Merging
author_association: COLLABORATOR