Repository: skypilot-org/skypilot
Pull request #3505 (2024-05-02 08:13:05)

[Core] Fix label validation for Kubernetes, add unit tests
Fixes label validation for Kubernetes, where it was not accepting the `kueue.x-k8s.io/queue-name` label as valid. Also added unit tests to make sure validators are tested.

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Label smoke tests - `pytest tests/test_smoke.py::test_task_labels_aws --aws, pytest tests/test_smoke.py::test_task_labels_gcp --gcp, pytest tests/test_smoke.py::test_task_labels_kubernetes --kubernetes`
- [x] New unit tests - `pytest tests/unit_tests/test_resources.py`


Patch:
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index 0d8e89835..9777a2894 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -403,11 +403,17 @@ class Kubernetes(clouds.Cloud):
     @classmethod
     def is_label_valid(cls, label_key: str,
                        label_value: str) -> Tuple[bool, Optional[str]]:
+        # Kubernetes labels can be of the format <domain>/<key>: <value>
         key_regex = re.compile(
-            r'^(?:[a-z0-9]([-a-z0-9]*[a-z0-9])?\/)?[a-z0-9]([-a-z0-9_.]{0,61}'
-            r'[a-z0-9])?$')
+            # Look-ahead to ensure proper domain formatting up to a slash
+            r'^(?:(?=[a-z0-9]([-a-z0-9.]*[a-z0-9])?\/)'
+            # Match domain: starts and ends with alphanum up to 253 chars
+            # including a slash in the domain.
+            r'[a-z0-9]([-a-z0-9.]{0,251}[a-z0-9])?\/)?'
+            # Match key: starts and ends with alphanum, upto to 63 chars.
+            r'[a-z0-9]([-a-z0-9_.]{0,61}[a-z0-9])?$')
         value_regex = re.compile(
-            r'^[a-zA-Z0-9]([-a-zA-Z0-9_.]{0,61}[a-zA-Z0-9])?$')
+            r'^([a-zA-Z0-9]([-a-zA-Z0-9_.]{0,61}[a-zA-Z0-9])?)?$')
         key_valid = bool(key_regex.match(label_key))
         value_valid = bool(value_regex.match(label_value))
         error_msg = None


Test patch:
diff --git a/tests/unit_tests/test_resources.py b/tests/unit_tests/test_resources.py
index fb2825e76..f9e3ad516 100644
--- a/tests/unit_tests/test_resources.py
+++ b/tests/unit_tests/test_resources.py
@@ -1,7 +1,23 @@
+from typing import Dict
 from unittest.mock import Mock
 
+import pytest
+
+from sky import clouds
 from sky.resources import Resources
 
+GLOBAL_VALID_LABELS = {
+    'plaintext': 'plainvalue',
+    'numbers123': '123',
+}
+
+GLOBAL_INVALID_LABELS = {
+    'l' * 129: 'value',  # Long key
+    'key': 'v' * 257,  # Long value
+    'spaces in label': 'spaces in value',
+    '': 'emptykey',
+}
+
 
 def test_get_reservations_available_resources():
     mock = Mock()
@@ -11,3 +27,62 @@ def test_get_reservations_available_resources():
     r.get_reservations_available_resources()
     mock.get_reservations_available_resources.assert_called_once_with(
         "instance_type", "region", "zone", set())
+
+
+def _run_label_test(allowed_labels: Dict[str, str],
+                    invalid_labels: Dict[str, str], cloud: clouds.Cloud):
+    """Run a test for labels with the given allowed and invalid labels."""
+    r_allowed = Resources(cloud=cloud, labels=allowed_labels)  # Should pass
+    assert r_allowed.labels == allowed_labels, ('Allowed labels '
+                                                'should be the same')
+
+    # Check for each invalid label
+    for invalid_label, value in invalid_labels.items():
+        l = {invalid_label: value}
+        with pytest.raises(ValueError):
+            _ = Resources(cloud=cloud, labels=l)
+            assert False, (f'Resources were initialized with '
+                           f'invalid label {invalid_label}={value}')
+
+
+def test_gcp_labels_resources():
+    allowed_labels = {
+        **GLOBAL_VALID_LABELS,
+    }
+    invalid_labels = {
+        **GLOBAL_INVALID_LABELS,
+        'domain/key': 'value',
+        '1numericstart': 'value',
+    }
+    cloud = clouds.GCP()
+    _run_label_test(allowed_labels, invalid_labels, cloud)
+
+
+def test_aws_labels_resources():
+    allowed_labels = {
+        **GLOBAL_VALID_LABELS,
+    }
+    invalid_labels = {
+        **GLOBAL_INVALID_LABELS,
+        'aws:cannotstartwithaws': 'value',
+    }
+    cloud = clouds.AWS()
+    _run_label_test(allowed_labels, invalid_labels, cloud)
+
+
+def test_kubernetes_labels_resources():
+    allowed_labels = {
+        **GLOBAL_VALID_LABELS,
+        'kueue.x-k8s.io/queue-name': 'queue',
+        'a' * 253 + '/' + 'k' * 63: 'v' *
+                                    63,  # upto 253 chars in domain, 63 in key
+        'mylabel': '',  # empty label values are allowed by k8s
+        '1numericstart': 'value',  # numeric start is allowed by k8s
+    }
+    invalid_labels = {
+        **GLOBAL_INVALID_LABELS,
+        'a' * 254 + '/' + 'k' * 64: 'v' *
+                                    63,  # exceed 253 chars in domain, 63 in key
+    }
+    cloud = clouds.Kubernetes()
+    _run_label_test(allowed_labels, invalid_labels, cloud)