Respect user provided label_selector preferred instance_type

Author: benieric

src/hyperpod_cli/commands/job.py

@@ -674,8 +674,12 @@ def start_job(
 
             label_selector = config["cluster"]["cluster_config"].setdefault("label_selector",{})
             required_labels = label_selector.setdefault("required", {})
+            preferred_labels = label_selector.setdefault("preferred", {})
 
-            if not required_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY):
+            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] = (
                     [str(instance_type)]
                 )

src/hyperpod_cli/validators/job_validator.py

@@ -187,8 +187,12 @@ def validate_yaml_content(data):
 
     label_selector = cluster_config_fields.setdefault("label_selector",{})
     required_labels = label_selector.setdefault("required", {})
+    preferred_labels = label_selector.setdefault("preferred", {})
 
-    if not required_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY):
+    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] = (
             [str(instance_type)]
         )

test/unit_tests/test_job.py

@@ -575,6 +575,86 @@ def capture_yaml_dump(config, *args, **kwargs):
         print(f"Output: {result.output}")
         if result.exception:
             print(f"Exception: {result.exception}")
+            
+    @mock.patch('subprocess.run')
+    @mock.patch("yaml.dump")
+    @mock.patch("os.path.exists", return_value=True)
+    @mock.patch("os.remove", return_value=None)
+    @mock.patch("hyperpod_cli.utils.get_cluster_console_url")
+    @mock.patch("hyperpod_cli.clients.kubernetes_client.KubernetesClient.__new__")
+    @mock.patch("hyperpod_cli.commands.job.JobValidator")
+    @mock.patch("boto3.Session")
+    def test_start_job_label_selector_preferred_instance_type(
+        self,
+        mock_boto3,
+        mock_validator_cls,
+        mock_kubernetes_client,
+        mock_get_console_link,
+        mock_remove,
+        mock_exists,
+        mock_yaml_dump,
+        mock_subprocess_run,
+    ):
+        # Setup mocks
+        mock_validator = mock_validator_cls.return_value
+        mock_validator.validate_aws_credential.return_value = True
+        mock_kubernetes_client.get_current_context_namespace.return_value = "kubeflow"
+        mock_get_console_link.return_value = "test-console-link"
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=['some_command'],
+            returncode=0,
+            stdout='Command executed successfully',
+            stderr=''
+        )
+
+        expected_default_label_selector_config = {
+            "preferred": {"beta.kubernetes.io/instance-type": ["ml.c5.xlarge"]},
+        }
+
+        # Capture the yaml.dump calls to inspect the config
+        configs_dumped = []
+        def capture_yaml_dump(config, *args, **kwargs):
+            configs_dumped.append(config)
+            print(f"Dumped config: {config}")
+            return None
+        mock_yaml_dump.side_effect = capture_yaml_dump
+
+        # Run the command
+        result = self.runner.invoke(
+            start_job,
+            [
+                "--job-name", "test-job",
+                "--instance-type", "ml.c5.xlarge",
+                "--image", "pytorch:1.9.0-cuda11.1-cudnn8-runtime",
+                "--node-count", "2",
+                "--entry-script", "/opt/train/src/train.py",
+                "--label_selector", 
+                '{"preferred": {"beta.kubernetes.io/instance-type": ["ml.c5.xlarge"]}}',
+            ],
+            catch_exceptions=False
+        )
+
+        # Verify the command executed successfully
+        self.assertEqual(result.exit_code, 0)
+
+        # Get the config that was generated
+        self.assertTrue(len(configs_dumped) > 0, "No config was generated")
+        config = configs_dumped[0]  # Get the first config that was dumped
+
+        # Verify label_selector configuration
+        self.assertIn('cluster', config)
+        self.assertIn('cluster_config', config['cluster'])
+        self.assertIn('label_selector', config['cluster']['cluster_config'])
+
+        self.assertEqual(
+            config['cluster']['cluster_config']['label_selector'],
+            expected_default_label_selector_config
+        )
+
+        print(f"Exit code: {result.exit_code}")
+        print(f"Output: {result.output}")
+        if result.exception:
+            print(f"Exception: {result.exception}")
 
     @mock.patch('subprocess.run')
     @mock.patch("yaml.dump")

test/unit_tests/validators/test_job_validator.py

@@ -1086,6 +1086,39 @@ def test_validate_yaml_content_valid(self):
         }
         result = validate_yaml_content(mock_data)
         self.assertTrue(result)
+        
+    def test_validate_yaml_content_preferred_instance_type_label(self):
+        expected_label_selector = {
+            "preferred": {
+                "beta.kubernetes.io/instance-type": [
+                    "ml.g5.xlarge"
+                ]
+            }
+        }
+        
+        # Respect user provided label_selector
+        mock_data = {
+            "cluster": {
+                "cluster_type": "k8s",
+                "instance_type": "ml.g5.xlarge",
+                "cluster_config": {
+                    "scheduler": "SageMaker",
+                    "label_selector": {
+                        "preferred": {
+                            "beta.kubernetes.io/instance-type": [
+                                "ml.g5.xlarge"
+                            ]
+                        }
+                    }
+                },
+            }
+        }
+        
+        result = validate_yaml_content(mock_data)
+        self.assertTrue(result)
+        self.assertEqual(
+            mock_data["cluster"]["cluster_config"]["label_selector"], expected_label_selector
+        )
 
     def test_validate_yaml_content_required_instance_type_label(self):
         expected_label_selector = {