Updated to handle YAML arrays in config file (#190)

Author: nargokul

src/sagemaker/hyperpod/cli/init_utils.py

@@ -226,7 +226,7 @@ def _parse_json_flag(ctx, param, value):
                     try:
                         # Remove brackets and split by comma
                         inner = value.strip()[1:-1]
-                        items = [item.strip() for item in inner.split(',')]
+                        items = [item.strip().strip('"').strip("'") for item in inner.split(',')]
                         return items
                     except:
                         pass
@@ -458,6 +458,15 @@ def save_config_yaml(prefill: dict, comment_map: dict, directory: str):
                     if vol.get('read_only') is not None:
                         f.write(f"    read_only: {vol.get('read_only')}\n")
                 f.write("\n")
+            elif isinstance(val, list):
+                # Handle arrays in YAML format
+                if val:
+                    f.write(f"{key}:\n")
+                    for item in val:
+                        f.write(f"  - {item}\n")
+                else:
+                    f.write(f"{key}: []\n")
+                f.write("\n")
             else:
                 # Handle simple values
                 val = "" if val is None else val
@@ -780,7 +789,7 @@ def build_config_from_schema(template: str, version: str, model_config=None, exi
                                 if val_stripped.startswith('[') and val_stripped.endswith(']'):
                                     try:
                                         inner = val_stripped[1:-1]
-                                        val = [item.strip() for item in inner.split(',')]
+                                        val = [item.strip().strip('"').strip("'") for item in inner.split(',')]
                                     except:
                                         pass
 

src/sagemaker/hyperpod/cluster_management/hp_cluster_stack.py

@@ -12,9 +12,6 @@
 
 from sagemaker.hyperpod import create_boto3_client
 
-CLUSTER_CREATION_TEMPLATE_FILE_NAME = "v1_0/main-stack-eks-based-cfn-template.yaml"
-CLUSTER_STACK_TEMPLATE_PACKAGE_NAME = "hyperpod_cluster_stack_template"
-
 CAPABILITIES_FOR_STACK_CREATION = [
 'CAPABILITY_IAM',
 'CAPABILITY_NAMED_IAM'
@@ -31,17 +28,25 @@ class HpClusterStack(ClusterStackBase):
         None,
         description="CloudFormation stack name set after stack creation"
     )
+    
+    def __init__(self, **data):
+        # Convert array values to JSON strings
+        for key, value in data.items():
+            if isinstance(value, list):
+                data[key] = json.dumps(value)
+        super().__init__(**data)
 
     @staticmethod
     def get_template() -> str:
-        s3 = create_boto3_client('s3')
-        response = s3.get_object(
-            Bucket='sagemaker-hyperpod-cluster-stack-bucket',
-            Key='1.0/main-stack-eks-based-cfn-template.yaml'
-        )
-        yaml_content = response['Body'].read().decode('utf-8')
-        yaml_data = yaml.safe_load(yaml_content)
-        return json.dumps(yaml_data, indent=2, ensure_ascii=False)
+        try:
+            template_content = importlib.resources.read_text(
+                'hyperpod_cluster_stack_template', 
+                'creation_template.yaml'
+            )
+            yaml_data = yaml.safe_load(template_content)
+            return json.dumps(yaml_data, indent=2, ensure_ascii=False)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load template from package: {e}")
 
     def create(self,
                region: Optional[str] = None) -> str:
@@ -56,7 +61,7 @@ def create(self,
         # Get account ID and create bucket name
         bucket_name = f"sagemaker-hyperpod-cluster-stack-bucket"
         template_key = f"1.0/main-stack-eks-based-cfn-template.yaml"
-        
+
         try:
             # Use TemplateURL for large templates (>51KB)
             template_url = f"https://{bucket_name}.s3.amazonaws.com/{template_key}"
@@ -65,10 +70,7 @@ def create(self,
                 StackName=stack_name,
                 TemplateURL=template_url,
                 Parameters=parameters,
-                Tags=self.tags or [{
-                        'Key': 'Environment',
-                        'Value': 'Development'
-                    }],
+                Tags=self._parse_tags(),
                 Capabilities=CAPABILITIES_FOR_STACK_CREATION
             )
 
@@ -137,6 +139,13 @@ def _create_parameters(self) -> List[Dict[str, str]]:
                     })
         return parameters
 
+    def _parse_tags(self) -> List[Dict[str, str]]:
+        """Parse tags field and return proper CloudFormation tags format."""
+        try:
+            return json.loads(self.tags) if self.tags else []
+        except (json.JSONDecodeError, TypeError):
+            return []
+
     def _convert_nested_keys(self, obj: Any) -> Any:
         """Convert nested JSON keys from snake_case to PascalCase."""
         if isinstance(obj, dict):

test/unit_tests/cli/test_init_utils.py

@@ -928,4 +928,45 @@ def test_process_cfn_template_content_preserves_template_structure(self, mock_cl
         assert "CloudFormation Template" in result
         assert original_content in result
         assert "OtherParam" in result
-        assert mock_template in result
+        assert mock_template in result
+
+
+class TestProcessCfnTemplateContentUpdated:
+    """Test updated _process_cfn_template_content function"""
+    
+    @patch('sagemaker.hyperpod.cli.init_utils.HpClusterStack.get_template')
+    def test_process_cfn_template_content_uses_package_template(self, mock_get_template):
+        """Test that _process_cfn_template_content uses HpClusterStack.get_template from package"""
+        original_content = "Original template content"
+        mock_template = '{"Parameters": {"TestParam": {"Type": "String"}}}'
+        mock_get_template.return_value = mock_template
+        
+        result = _process_cfn_template_content(original_content)
+        
+        # Verify get_template was called (not S3)
+        mock_get_template.assert_called_once()
+        
+        # Verify content structure
+        assert "CloudFormation Template:" in result
+        assert mock_template in result
+        assert original_content in result
+    
+    @patch('sagemaker.hyperpod.cli.init_utils.HpClusterStack.get_template')
+    @patch('sagemaker.hyperpod.cli.init_utils.click.secho')
+    def test_process_cfn_template_content_handles_package_error(self, mock_secho, mock_get_template):
+        """Test error handling when package template fails to load"""
+        original_content = "Original template content"
+        mock_get_template.side_effect = RuntimeError("Failed to load template from package")
+        
+        result = _process_cfn_template_content(original_content)
+        
+        # Verify error was logged
+        mock_secho.assert_called_once_with(
+            "⚠️ Failed to generate CloudFormation template: Failed to load template from package", 
+            fg="red"
+        )
+        
+        # Verify fallback behavior
+        assert "CloudFormation Template:" in result
+        assert original_content in result
+        assert result.count("CloudFormation Template:") == 1

test/unit_tests/cluster_management/test_hp_cluster_stack.py

@@ -4,8 +4,7 @@
 from botocore.exceptions import ClientError
 import boto3
 
-from sagemaker.hyperpod.cluster_management.hp_cluster_stack import HpClusterStack, CLUSTER_STACK_TEMPLATE_PACKAGE_NAME, CLUSTER_CREATION_TEMPLATE_FILE_NAME
-
+from sagemaker.hyperpod.cluster_management.hp_cluster_stack import HpClusterStack
 
 class TestHpClusterStack(unittest.TestCase):
     @patch('uuid.uuid4')
@@ -392,9 +391,134 @@ def test_create_parameters_preserves_other_fields(self):
 
         # Should have the other fields
         param_keys = [p['ParameterKey'] for p in other_params]
-        self.assertIn('HyperpodClusterName', param_keys)
+        self.assertIn('HyperPodClusterName', param_keys)
         self.assertIn('CreateVPCStack', param_keys)
 
         # Verify boolean conversion
         vpc_param = next(p for p in other_params if p['ParameterKey'] == 'CreateVPCStack')
-        self.assertEqual(vpc_param['ParameterValue'], 'true')
+        self.assertEqual(vpc_param['ParameterValue'], 'true')
+
+class TestHpClusterStackInit(unittest.TestCase):
+    """Test HpClusterStack __init__ method array conversion"""
+    
+    def test_init_converts_arrays_to_json_strings(self):
+        """Test that __init__ converts array values to JSON strings"""
+        data = {
+            'tags': [{'Key': 'Environment', 'Value': 'Test'}],
+            'availability_zone_ids': ['us-east-1a', 'us-east-1b'],
+            'hyperpod_cluster_name': 'test-cluster',
+            'storage_capacity': 1200
+        }
+        
+        stack = HpClusterStack(**data)
+        
+        # Arrays should be converted to JSON strings
+        self.assertEqual(stack.tags, '[{"Key": "Environment", "Value": "Test"}]')
+        self.assertEqual(stack.availability_zone_ids, '["us-east-1a", "us-east-1b"]')
+        
+        # Other types should remain unchanged
+        self.assertEqual(stack.hyperpod_cluster_name, 'test-cluster')
+        self.assertEqual(stack.storage_capacity, 1200)
+    
+    def test_init_handles_empty_arrays(self):
+        """Test that empty arrays are converted to empty JSON arrays"""
+        data = {'tags': []}
+        
+        stack = HpClusterStack(**data)
+        
+        self.assertEqual(stack.tags, '[]')
+    
+    def test_init_handles_no_arrays(self):
+        """Test that __init__ works normally when no arrays are present"""
+        data = {
+            'hyperpod_cluster_name': 'test-cluster',
+            'stage': 'gamma'
+        }
+        
+        stack = HpClusterStack(**data)
+        
+        self.assertEqual(stack.hyperpod_cluster_name, 'test-cluster')
+        self.assertEqual(stack.stage, 'gamma')
+
+
+class TestHpClusterStackParseTags(unittest.TestCase):
+    """Test HpClusterStack _parse_tags method"""
+    
+    def test_parse_tags_valid_json_array(self):
+        """Test parsing valid JSON array of tags"""
+        tags_json = '[{"Key": "Environment", "Value": "Test"}, {"Key": "Project", "Value": "HyperPod"}]'
+        stack = HpClusterStack(tags=tags_json)
+        
+        result = stack._parse_tags()
+        
+        expected = [
+            {"Key": "Environment", "Value": "Test"},
+            {"Key": "Project", "Value": "HyperPod"}
+        ]
+        self.assertEqual(result, expected)
+    
+    def test_parse_tags_empty_string(self):
+        """Test parsing empty tags string returns empty list"""
+        stack = HpClusterStack(tags="")
+        
+        result = stack._parse_tags()
+        
+        self.assertEqual(result, [])
+    
+    def test_parse_tags_none_value(self):
+        """Test parsing None tags returns empty list"""
+        stack = HpClusterStack(tags=None)
+        
+        result = stack._parse_tags()
+        
+        self.assertEqual(result, [])
+    
+    def test_parse_tags_invalid_json(self):
+        """Test parsing invalid JSON returns empty list"""
+        stack = HpClusterStack(tags="invalid json")
+        
+        result = stack._parse_tags()
+        
+        self.assertEqual(result, [])
+    
+    def test_parse_tags_empty_json_array(self):
+        """Test parsing empty JSON array returns empty list"""
+        stack = HpClusterStack(tags="[]")
+        
+        result = stack._parse_tags()
+        
+        self.assertEqual(result, [])
+
+
+class TestHpClusterStackGetTemplate(unittest.TestCase):
+    """Test HpClusterStack get_template method using package instead of S3"""
+    
+    @patch('sagemaker.hyperpod.cluster_management.hp_cluster_stack.importlib.resources.read_text')
+    @patch('sagemaker.hyperpod.cluster_management.hp_cluster_stack.yaml.safe_load')
+    def test_get_template_from_package(self, mock_yaml_load, mock_read_text):
+        """Test get_template reads from package instead of S3"""
+        mock_yaml_content = "Parameters:\n  TestParam:\n    Type: String"
+        mock_read_text.return_value = mock_yaml_content
+        
+        mock_yaml_data = {"Parameters": {"TestParam": {"Type": "String"}}}
+        mock_yaml_load.return_value = mock_yaml_data
+        
+        result = HpClusterStack.get_template()
+        
+        # Verify package resource was read
+        mock_read_text.assert_called_once_with('hyperpod_cluster_stack_template', 'creation_template.yaml')
+        mock_yaml_load.assert_called_once_with(mock_yaml_content)
+        
+        # Verify JSON output
+        expected_json = json.dumps(mock_yaml_data, indent=2, ensure_ascii=False)
+        self.assertEqual(result, expected_json)
+    
+    @patch('sagemaker.hyperpod.cluster_management.hp_cluster_stack.importlib.resources.read_text')
+    def test_get_template_handles_package_error(self, mock_read_text):
+        """Test get_template handles package read errors"""
+        mock_read_text.side_effect = FileNotFoundError("Template not found")
+        
+        with self.assertRaises(RuntimeError) as context:
+            HpClusterStack.get_template()
+        
+        self.assertIn("Failed to load template from package", str(context.exception))