match skypilot labeler logic

skypilot-org · May 16, 2024 · 7d4d6bb · 7d4d6bb
1 parent 75e4396
commit 7d4d6bb
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/sky/provision/kubernetes/utils.py b/sky/provision/kubernetes/utils.py
@@ -206,8 +206,7 @@ def get_label_key(cls) -> str:
     @classmethod
     def get_accelerator_from_label_value(cls, value: str) -> str:
         """Searches against a canonical list of NVIDIA GPUs and pattern
-        matches the canonical GPU name against the GFD label. Taken from
-        sky/utils/kubernetes/k8s_gpu_labeler_setup.yaml
+        matches the canonical GPU name against the GFD label.
         """
         canonical_gpu_names = [
             'A100-80GB', 'A100', 'A10G', 'H100', 'K80', 'M60', 'T4g', 'T4',
@@ -224,7 +223,10 @@ def get_accelerator_from_label_value(cls, value: str) -> str:
         # If we didn't find a canonical name:
         # 1. remove 'NVIDIA-' (e.g., 'NVIDIA-RTX-A6000' -> 'RTX-A6000')
         # 2. remove 'GEFORCE-' (e.g., 'NVIDIA-GEFORCE-RTX-3070' -> 'RTX-3070')
-        return value.upper().replace('NVIDIA-', '').replace('GEFORCE-', '')
+        # 3. remove 'RTX-' (e.g. 'RTX-6000' -> 'RTX6000')
+        # Same logic, but uppercased, as the Skypilot labeler job found in
+        # sky/utils/kubernetes/k8s_gpu_labeler_setup.yaml
+        return value.upper().replace('NVIDIA-', '').replace('GEFORCE-', '').replace('RTX-', 'RTX')
 
 
 class KarpenterLabelFormatter(SkyPilotLabelFormatter):