[k8s] Add label validation during sky check (#2653)

* Add label validation

* early return

* lint

* update docs

* fix test

* update monkey patching

* fixes

docs/source/reference/kubernetes/kubernetes-setup.rst

@@ -167,10 +167,10 @@ Please follow their respective guides to deploy your Kubernetes cluster.
 
 Setting up GPU support
 ~~~~~~~~~~~~~~~~~~~~~~
-If your Kubernetes cluster has Nvidia GPUs, make sure you have the Nvidia
-device plugin installed (i.e., ``nvidia.com/gpu`` resource is available on each node).
-Additionally, you will need to label each node in your cluster with the GPU type.
-For example, a node with v100 GPUs must have a label :code:`skypilot.co/accelerators: v100`.
+If your Kubernetes cluster has Nvidia GPUs, ensure that:
+
+1. The Nvidia device plugin is installed (i.e., ``nvidia.com/gpu`` resource is available on each node).
+2. Each node in your cluster is labelled with the GPU type. This labelling can be done by adding a label of the format ``skypilot.co/accelerators: <gpu_name>``, where the ``<gpu_name>`` is the lowercase name of the GPU. For example, a node with V100 GPUs must have a label :code:`skypilot.co/accelerators: v100`.
 
 We provide a convenience script that automatically detects GPU types and labels each node. You can run it with:
 
@@ -185,6 +185,9 @@ We provide a convenience script that automatically detects GPU types and labels
  You can check if nodes have been labeled by running `kubectl describe nodes` and looking for labels of the format `skypilot.co/accelerators: <gpu_name>`.
 
 
+.. note::
+ GPU labels are case-sensitive. Ensure that the GPU name is lowercase if you are using the ``skypilot.co/accelerators`` label.
+
 .. note::
  GPU labelling is not required on GKE clusters - SkyPilot will automatically use GKE provided labels. However, you will still need to install `drivers <https://cloud.google.com/kubernetes-engine/docs/how-to/gpus#installing_drivers>`_.
 

sky/utils/kubernetes_utils.py

@@ -60,6 +60,20 @@ def get_accelerator_from_label_value(cls, value: str) -> str:
  """Given a label value, returns the GPU type"""
  raise NotImplementedError
 
+ @classmethod
+ def validate_label_value(cls, value: str) -> Tuple[bool, str]:
+ """Validates if the specified label value is correct.
+
+ Used to check if the labelling on the cluster is correct and
+ preemptively raise an error if it is not.
+
+ Returns:
+ bool: True if the label value is valid, False otherwise.
+ str: Error message if the label value is invalid, None otherwise.
+ """
+ del value
+ return True, ''
+
 
 def get_gke_accelerator_name(accelerator: str) -> str:
  """Returns the accelerator name for GKE clusters
@@ -97,6 +111,14 @@ def get_label_value(cls, accelerator: str) -> str:
  def get_accelerator_from_label_value(cls, value: str) -> str:
  return value.upper()
 
+ @classmethod
+ def validate_label_value(cls, value: str) -> Tuple[bool, str]:
+ """Values must be all lowercase for the SkyPilot formatter."""
+ is_valid = value == value.lower()
+ return is_valid, (f'Label value {value!r} must be lowercase if using '
+ f'the {cls.get_label_key()} label.'
+ if not is_valid else '')
+
 
 class CoreWeaveLabelFormatter(GPULabelFormatter):
  """CoreWeave label formatter
@@ -157,31 +179,32 @@ def get_accelerator_from_label_value(cls, value: str) -> str:
 
 
 def detect_gpu_label_formatter(
-) -> Tuple[Optional[GPULabelFormatter], List[Tuple[str, str]]]:
+) -> Tuple[Optional[GPULabelFormatter], Dict[str, List[Tuple[str, str]]]]:
  """Detects the GPU label formatter for the Kubernetes cluster
 
  Returns:
  GPULabelFormatter: The GPU label formatter for the cluster, if found.
- List[Tuple[str, str]]: The set of labels and values across all nodes.
+ Dict[str, List[Tuple[str, str]]]: A mapping of nodes and the list of
+ labels on each node. E.g., {'node1': [('label1', 'value1')]}
  """
  # Get all labels across all nodes
- node_labels: List[Tuple[str, str]] = []
+ node_labels: Dict[str, List[Tuple[str, str]]] = {}
  nodes = get_kubernetes_nodes()
  for node in nodes:
- node_labels.extend(node.metadata.labels.items())
+ node_labels[node.metadata.name] = []
+ for label, value in node.metadata.labels.items():
+ node_labels[node.metadata.name].append((label, value))
 
  label_formatter = None
 
  # Check if the node labels contain any of the GPU label prefixes
  for lf in LABEL_FORMATTER_REGISTRY:
  label_key = lf.get_label_key()
- for label, _ in node_labels:
- if label.startswith(label_key):
- label_formatter = lf()
- break
-
- if label_formatter is not None:
- break
+ for _, label_list in node_labels.items():
+ for label, _ in label_list:
+ if label.startswith(label_key):
+ label_formatter = lf()
+ return label_formatter, node_labels
 
  return label_formatter, node_labels
 
@@ -343,6 +366,17 @@ def get_gpu_label_key_value(acc_type: str, check_mode=False) -> Tuple[str, str]:
  'the documentation on how to set up node labels.'
  f'{suffix}')
  if label_formatter is not None:
+ # Validate the label value on all nodes labels to ensure they are
+ # correctly setup and will behave as expected.
+ for node_name, label_list in node_labels.items():
+ for label, value in label_list:
+ if label == label_formatter.get_label_key():
+ is_valid, reason = label_formatter.validate_label_value(
+ value)
+ if not is_valid:
+ raise exceptions.ResourcesUnavailableError(
+ f'Node {node_name!r} in Kubernetes cluster has '
+ f'invalid GPU label: {label}={value}. {reason}')
  if check_mode:
  # If check mode is enabled and we reached so far, we can
  # conclude that the cluster is setup correctly and return.
@@ -355,17 +389,22 @@ def get_gpu_label_key_value(acc_type: str, check_mode=False) -> Tuple[str, str]:
  # node. It does not (and should not) check if the resource
  # quantity is available since that is dynamic and can change
  # during scheduling.
- for label, value in node_labels:
- if label == k8s_acc_label_key and value == k8s_acc_label_value:
- # If a node is found, we can break out of the loop
- # and proceed to deploy.
- return k8s_acc_label_key, k8s_acc_label_value
+ for node_name, label_list in node_labels.items():
+ for label, value in label_list:
+ if (label == k8s_acc_label_key and
+ value == k8s_acc_label_value):
+ # If a node is found, we can break out of the loop
+ # and proceed to deploy.
+ return k8s_acc_label_key, k8s_acc_label_value
  # If no node is found with the requested acc_type, raise error
  with ux_utils.print_exception_no_traceback():
  suffix = ''
  if env_options.Options.SHOW_DEBUG_INFO.get():
+ all_labels = []
+ for node_name, label_list in node_labels.items():
+ all_labels.extend(label_list)
  gpus_available = set(
- v for k, v in node_labels if k == k8s_acc_label_key)
+ v for k, v in all_labels if k == k8s_acc_label_key)
  suffix = f' Available GPUs on the cluster: {gpus_available}'
  raise exceptions.ResourcesUnavailableError(
  'Could not find any node in the Kubernetes cluster '

tests/common.py

@@ -55,7 +55,7 @@ def _get_az_mappings(_):
  # the cluster to detect available cluster resources.
  monkeypatch.setattr(
  'sky.utils.kubernetes_utils.detect_gpu_label_formatter',
- lambda *_args, **_kwargs: [kubernetes_utils.SkyPilotLabelFormatter, []])
+ lambda *_args, **_kwargs: [kubernetes_utils.SkyPilotLabelFormatter, {}])
  monkeypatch.setattr('sky.utils.kubernetes_utils.detect_gpu_resource',
  lambda *_args, **_kwargs: [True, []])
  monkeypatch.setattr('sky.utils.kubernetes_utils.check_instance_fits',