change: autofill required instance-type label

benieric · benieric · commit 4fe9a3799af4 · 2025-03-13T15:15:19.000-07:00
diff --git a/src/hyperpod_cli/commands/job.py b/src/hyperpod_cli/commands/job.py
@@ -32,6 +32,7 @@
     HYPERPOD_KUBERNETES_JOB_PREFIX,
     HYPERPOD_MAX_RETRY_ANNOTATION_KEY,
     HYPERPOD_NAMESPACE_PREFIX,
+    KUBERNETES_INSTANCE_TYPE_LABEL_KEY,
     KUEUE_JOB_UID_LABEL_KEY,
     KUEUE_QUEUE_NAME_LABEL_KEY,
     KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY,
@@ -671,6 +672,15 @@ def start_job(
                     NODE_AFFINITY_DICT
                 )
 
+            label_selector = config["cluster"].setdefault("label_selector",{})
+            required_labels = label_selector.setdefault("required", {})
+
+            if not required_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY):
+                required_labels[KUBERNETES_INSTANCE_TYPE_LABEL_KEY] = (
+                    [str(instance_type)]
+                )
+
+
             if auto_resume:
                 # Set max_retry default to 1
                 if max_retry is None:
diff --git a/src/hyperpod_cli/constants/command_constants.py b/src/hyperpod_cli/constants/command_constants.py
@@ -26,6 +26,7 @@
         "sagemaker.amazonaws.com/deep-health-check-status": ["Passed"],
     },
 }
+KUBERNETES_INSTANCE_TYPE_LABEL_KEY = "beta.kubernetes.io/instance-type"
 KUEUE_QUEUE_NAME_LABEL_KEY = "kueue.x-k8s.io/queue-name"
 KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY = "kueue.x-k8s.io/priority-class"
 KUEUE_JOB_UID_LABEL_KEY = "kueue.x-k8s.io/job-uid"
diff --git a/src/hyperpod_cli/validators/job_validator.py b/src/hyperpod_cli/validators/job_validator.py
@@ -24,6 +24,7 @@
     KUEUE_QUEUE_NAME_LABEL_KEY,
     HYPERPOD_AUTO_RESUME_ANNOTATION_KEY,
     HYPERPOD_MAX_RETRY_ANNOTATION_KEY,
+    KUBERNETES_INSTANCE_TYPE_LABEL_KEY,
     SchedulerType
 )
 from hyperpod_cli.constants.hyperpod_instance_types import (
@@ -184,6 +185,19 @@ def validate_yaml_content(data):
     if custom_labels is not None:
         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", {})
+
+    if not required_labels.get(KUBERNETES_INSTANCE_TYPE_LABEL_KEY):
+        required_labels[KUBERNETES_INSTANCE_TYPE_LABEL_KEY] = (
+            [str(instance_type)]
+        )
+    if instance_type not in required_labels[KUBERNETES_INSTANCE_TYPE_LABEL_KEY]:
+        logger.error(
+            f"Please ensure 'instance-type' in 'cluster' matches with 'instance-type' in 'label_selector.required.beta.kubernetes.io/instance-type' in config file"
+        )
+        return False
+
     auto_resume = False
     max_retry = None
     if annotations is not None:
diff --git a/test/integration_tests/data/basicJob.yaml b/test/integration_tests/data/basicJob.yaml
@@ -22,7 +22,7 @@ training_cfg:
  entry_script: /opt/pytorch-mnist/mnist.py
  script_args: []
  run:
-  name: ${JOB_NAME} # Current run name
+  name: hyperpod-benieric-test # Current run name
   nodes: 1 # Number of nodes to use for current training
   ntasks_per_node: 1 # Number of devices to use per node
 cluster:
@@ -40,6 +40,8 @@ cluster:
       required:
           sagemaker.amazonaws.com/node-health-status:
               - Schedulable
+          beta.kubernetes.io/instance-type:
+              - ml.c5.4xlarge
       preferred:
           sagemaker.amazonaws.com/deep-health-check-status:
               - Passed
diff --git a/test/unit_tests/test_job.py b/test/unit_tests/test_job.py
@@ -493,6 +493,90 @@ def test_start_job_with_cli_args(
             print(f"Exception: {result.exception}")
         self.assertEqual(result.exit_code, 0)
 
+    @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_default_label_selector_config(
+        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 = {
+            "required": {
+                "sagemaker.amazonaws.com/node-health-status": ["Schedulable"],
+                "beta.kubernetes.io/instance-type": ["ml.c5.xlarge"]
+            },
+            "preferred": {"sagemaker.amazonaws.com/deep-health-check-status": ["Passed"]},
+            "weights": [100],
+        }
+
+        # 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",
+            ],
+            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")
     @mock.patch("os.path.exists", return_value=True)
diff --git a/test/unit_tests/validators/test_job_validator.py b/test/unit_tests/validators/test_job_validator.py
@@ -1086,6 +1086,119 @@ def test_validate_yaml_content_valid(self):
         }
         result = validate_yaml_content(mock_data)
         self.assertTrue(result)
+    
+    def test_validate_yaml_content_valid_required_instance_type_label(self):
+        # Test auto-populate label selector - "beta.kubernetes.io/instance_type"
+        expected_label_selector = {
+            "required": {
+                "beta.kubernetes.io/instance-type": [
+                    "ml.g5.xlarge"
+                ]
+            }
+        }
+        
+        mock_data = {
+            "cluster": {
+                "cluster_type": "k8s",
+                "instance_type": "ml.g5.xlarge",
+                "cluster_config": {
+                    "scheduler": "SageMaker"
+                },
+            }
+        }
+        
+        result = validate_yaml_content(mock_data)
+        self.assertTrue(result)
+        self.assertEqual(
+            mock_data["cluster"]["cluster_config"]["label_selector"], expected_label_selector
+        )
+        
+        expected_label_selector = {
+            "required": {
+                "sagemaker.amazonaws.com/node-health-status": [
+                    "Schedulable"
+                ],
+                "beta.kubernetes.io/instance-type": [
+                    "ml.g5.xlarge"
+                ]
+            }
+        }
+        
+        mock_data = {
+            "cluster": {
+                "cluster_type": "k8s",
+                "instance_type": "ml.g5.xlarge",
+                "cluster_config": {
+                    "scheduler": "SageMaker",
+                    "label_selector": {
+                        "required": {
+                            "sagemaker.amazonaws.com/node-health-status": [
+                                "Schedulable"
+                            ]
+                        }
+                    }
+                },
+            }
+        }
+        
+        result = validate_yaml_content(mock_data)
+        self.assertTrue(result)
+        self.assertEqual(
+            mock_data["cluster"]["cluster_config"]["label_selector"], expected_label_selector
+        )
+        
+        expected_label_selector = {
+            "required": {
+                "beta.kubernetes.io/instance-type": [
+                    "ml.g5.xlarge"
+                ]
+            }
+        }
+
+        mock_data = {
+            "cluster": {
+                "cluster_type": "k8s",
+                "instance_type": "ml.g5.xlarge",
+                "cluster_config": {
+                    "scheduler": "SageMaker",
+                    "label_selector": {
+                        "required": {
+                            "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_invalid_required_instance_type_label(self):
+        # Test resepect user selection label selector
+        mock_data = {
+            "cluster": {
+                "cluster_type": "k8s",
+                "instance_type": "ml.g5.xlarge",
+                "cluster_config": {
+                    "scheduler": "SageMaker",
+                    "label_selector": {
+                        "required": {
+                            "beta.kubernetes.io/instance-type": [
+                                "ml.g5.2xlarge"
+                            ]
+                        }
+                    }
+                }
+            }
+        }
+
+        result = validate_yaml_content(mock_data)
+        self.assertFalse(result)
 
     def test_validate_yaml_content_error_no_cluster(
         self,

Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@`
`26`	`26`	`"sagemaker.amazonaws.com/deep-health-check-status": ["Passed"],`
`27`	`27`	`},`
`28`	`28`	`}`
	`29`	`+KUBERNETES_INSTANCE_TYPE_LABEL_KEY = "beta.kubernetes.io/instance-type"`
`29`	`30`	`KUEUE_QUEUE_NAME_LABEL_KEY = "kueue.x-k8s.io/queue-name"`
`30`	`31`	`KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY = "kueue.x-k8s.io/priority-class"`
`31`	`32`	`KUEUE_JOB_UID_LABEL_KEY = "kueue.x-k8s.io/job-uid"`