safer assignment of label

Author: benieric

src/hyperpod_cli/commands/job.py

@@ -673,14 +673,16 @@ def start_job(
                 )
 
             label_selector = config["cluster"]["cluster_config"].setdefault("label_selector",{})
-            required_labels = label_selector.setdefault("required", {})
-            preferred_labels = label_selector.setdefault("preferred", {})
+            required_labels = label_selector.get("required", {})
+            preferred_labels = label_selector.get("preferred", {})
 
             if (
                 not required_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY) and
                 not preferred_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY)
             ):
-                required_labels[KUBERNETES_INSTANCE_TYPE_LABEL_KEY] = (
+                if "required" not in label_selector:
+                    label_selector["required"] = {}
+                label_selector["required"][KUBERNETES_INSTANCE_TYPE_LABEL_KEY] = (
                     [str(instance_type)]
                 )
 

src/hyperpod_cli/validators/job_validator.py

@@ -186,13 +186,15 @@ def validate_yaml_content(data):
         queue_name = custom_labels.get(KUEUE_QUEUE_NAME_LABEL_KEY, None)
 
     label_selector = cluster_config_fields.setdefault("label_selector",{})
-    required_labels = label_selector.setdefault("required", {})
-    preferred_labels = label_selector.setdefault("preferred", {})
+    required_labels = label_selector.get("required", {})
+    preferred_labels = label_selector.get("preferred", {})
 
     if (
         not required_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY) and 
         not preferred_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY)
     ):
+        if "required" not in label_selector:
+            label_selector["required"] = {}
         required_labels[KUBERNETES_INSTANCE_TYPE_LABEL_KEY] = (
             [str(instance_type)]
         )

test/unit_tests/test_job.py

@@ -628,7 +628,7 @@ def capture_yaml_dump(config, *args, **kwargs):
                 "--image", "pytorch:1.9.0-cuda11.1-cudnn8-runtime",
                 "--node-count", "2",
                 "--entry-script", "/opt/train/src/train.py",
-                "--label_selector", 
+                "--label-selector", 
                 '{"preferred": {"beta.kubernetes.io/instance-type": ["ml.c5.xlarge"]}}',
             ],
             catch_exceptions=False