Cluster-stack template agnostic change (#245)

* decouple template from src code

* remove field validator from SDK pydantic model, fix minor parsing problem with list, update kubernetes_version type from str to float

* change type handler from class to module functions, change some public function to private, update unit tests

* cluster-stack template agnostic change

* update unit tests

* update integ test

* resolve circular import for cluster_stack

* resolve rebase merge conflict

* rename to_domain to to_config for cluster_stack

* increase timeout for endpoint integ test from 15min to 20min

Author: mollyheamazon

hyperpod-cluster-stack-template/hyperpod_cluster_stack_template/v1_0/model.py

@@ -1,5 +1,6 @@
 from pydantic import BaseModel, Field, field_validator
 from typing import Optional, Literal, List, Any, Union
+from sagemaker.hyperpod.common.utils import region_to_az_ids
 
 class ClusterStackBase(BaseModel):
     resource_name_prefix: Optional[str] = Field("hyp-eks-stack", description="Prefix to be used for all resources. A 4-digit UUID will be added to prefix during submission")
@@ -56,4 +57,77 @@ class ClusterStackBase(BaseModel):
     def validate_kubernetes_version(cls, v):
         if v is not None:
             return str(v)
-        return v
+        return v
+
+    def to_config(self, region: str = None):
+        """Convert CLI model to SDK configuration for cluster stack creation.
+        
+        Transforms the CLI model instance into a configuration dictionary that can be used
+        to instantiate the HpClusterStack SDK class. Applies necessary transformations
+        including AZ configuration, UUID generation, and field restructuring.
+        
+        Args:
+            region (str, optional): AWS region for AZ configuration. If provided,
+                automatically sets availability_zone_ids and fsx_availability_zone_id
+                when not already specified.
+        
+        Returns:
+            dict: Configuration dictionary ready for HpClusterStack instantiation.
+                Contains all transformed parameters with defaults applied.
+        
+        Example:
+            >>> cli_model = ClusterStackBase(hyperpod_cluster_name="my-cluster")
+            >>> config = cli_model.to_config(region="us-west-2")
+            >>> sdk_instance = HpClusterStack(**config)
+        """
+        import uuid
+        
+        # Convert model to dict and apply transformations
+        config = self.model_dump(exclude_none=True)
+        
+        # Prepare CFN arrays from numbered fields
+        instance_group_settings = []
+        rig_settings = []
+        for i in range(1, 21):
+            ig_key = f'instance_group_settings{i}'
+            rig_key = f'rig_settings{i}'
+            if ig_key in config:
+                instance_group_settings.append(config.pop(ig_key))
+            if rig_key in config:
+                rig_settings.append(config.pop(rig_key))
+        
+        # Add arrays to config
+        if instance_group_settings:
+            config['instance_group_settings'] = instance_group_settings
+        if rig_settings:
+            config['rig_settings'] = rig_settings
+        
+        # Add default AZ configuration if not provided
+        if region and (not config.get('availability_zone_ids') or not config.get('fsx_availability_zone_id')):
+            all_az_ids = region_to_az_ids(region)
+            default_az_config =     {
+                'availability_zone_ids': all_az_ids[:2],  # First 2 AZs
+                'fsx_availability_zone_id': all_az_ids[0]  # First AZ
+            }
+            if not config.get('availability_zone_ids'):
+                config['availability_zone_ids'] = default_az_config['availability_zone_ids']
+            if not config.get('fsx_availability_zone_id'):
+                config['fsx_availability_zone_id'] = default_az_config['fsx_availability_zone_id']
+        
+        # Append 4-digit UUID to resource_name_prefix
+        if config.get('resource_name_prefix'):
+            config['resource_name_prefix'] = f"{config['resource_name_prefix']}-{str(uuid.uuid4())[:4]}"
+        
+        # Set fixed defaults
+        defaults = {
+            'custom_bucket_name': 'sagemaker-hyperpod-cluster-stack-bucket',
+            'github_raw_url': 'https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/refs/heads/main/1.architectures/7.sagemaker-hyperpod-eks/LifecycleScripts/base-config/on_create.sh',
+            'helm_repo_url': 'https://github.com/aws/sagemaker-hyperpod-cli.git',
+            'helm_repo_path': 'helm_chart/HyperPodHelmChart'
+        }
+        
+        for key, default_value in defaults.items():
+            if key not in config:
+                config[key] = default_value
+        
+        return config

src/sagemaker/hyperpod/cli/commands/cluster_stack.py

@@ -18,6 +18,10 @@
 from sagemaker.hyperpod.common.telemetry.constants import Feature
 from sagemaker.hyperpod.common.utils import setup_logging
 from sagemaker.hyperpod.cli.utils import convert_datetimes
+from sagemaker.hyperpod.cli.init_utils import _filter_cli_metadata_fields
+from sagemaker.hyperpod.cli.init_utils import load_config
+from sagemaker.hyperpod.cli.constants.init_constants import TEMPLATES
+from pathlib import Path
 from sagemaker.hyperpod.cli.cluster_stack_utils import (
     StackNotFoundError,
     delete_stack_with_confirmation
@@ -67,86 +71,42 @@ def create_cluster_stack(config_file, region, debug):
           # Create with debug logging
           hyp create hyp-cluster cluster-config.yaml my-stack-name --debug
     """
-    create_cluster_stack_helper(config_file, region, debug)
-
-def create_cluster_stack_helper(config_file: str, region: Optional[str] = None, debug: bool = False) -> None:
-    """Helper function to create a HyperPod cluster stack.
-
-    **Parameters:**
-
-    .. list-table::
-       :header-rows: 1
-       :widths: 20 20 60
-
-       * - Parameter
-         - Type
-         - Description
-       * - config_file
-         - str
-         - Path to the YAML configuration file containing cluster stack settings
-       * - region
-         - str, optional
-         - AWS region where the cluster stack will be created
-       * - debug
-         - bool
-         - Enable debug logging for detailed error information
-
-    **Raises:**
-
-    ClickException: When cluster stack creation fails or configuration is invalid
-    """
     try:
         # Validate the config file path
         if not os.path.exists(config_file):
             logger.error(f"Config file not found: {config_file}")
             return
 
-        # Load the configuration from the YAML file
-        import yaml
-        import uuid
-        with open(config_file, 'r') as f:
-            config_data = yaml.safe_load(f)
-
-        # Filter out template and namespace fields
-        filtered_config = {}
-        for k, v in config_data.items():
-            if k not in ('template', 'namespace') and v is not None:
-                # Append 4-digit UUID to resource_name_prefix
-                if k == 'resource_name_prefix' and v:
-                    v = f"{v}-{str(uuid.uuid4())[:4]}"
-                filtered_config[k] = v
-
-        # Create the HpClusterStack object
-        # Ensure fixed defaults are always set
-        if 'custom_bucket_name' not in filtered_config:
-            filtered_config['custom_bucket_name'] = 'sagemaker-hyperpod-cluster-stack-bucket'
-        if 'github_raw_url' not in filtered_config:
-            filtered_config['github_raw_url'] = 'https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/refs/heads/main/1.architectures/7.sagemaker-hyperpod-eks/LifecycleScripts/base-config/on_create.sh'
-        if 'helm_repo_url' not in filtered_config:
-            filtered_config['helm_repo_url'] = 'https://github.com/aws/sagemaker-hyperpod-cli.git'
-        if 'helm_repo_path' not in filtered_config:
-            filtered_config['helm_repo_path'] = 'helm_chart/HyperPodHelmChart'
-        
-        cluster_stack = HpClusterStack(**filtered_config)
+        # Load config to get template and version
+
+        config_dir = Path(config_file).parent
+        data, template, version = load_config(config_dir)
+
+        # Get model from registry
+        registry = TEMPLATES[template]["registry"]
+        model_class = registry.get(str(version))
+
+        if model_class:
+            # Filter out CLI metadata fields
+            filtered_config = _filter_cli_metadata_fields(data)
 
-        # Log the configuration
-        logger.info("Creating HyperPod cluster stack with the following configuration:")
-        for key, value in filtered_config.items():
-            if value is not None:
-                logger.info(f"  {key}: {value}")
+            # Create model instance and domain
+            model_instance = model_class(**filtered_config)
+            config = model_instance.to_config(region=region)
 
-        # Create the cluster stack
-        stack_id = cluster_stack.create(region)
+            # Create the cluster stack
+            stack_id = HpClusterStack(**config).create(region)
 
-        logger.info(f"Stack creation initiated successfully with ID: {stack_id}")
-        logger.info("You can monitor the stack creation in the AWS CloudFormation console.")
+            logger.info(f"Stack creation initiated successfully with ID: {stack_id}")
+            logger.info("You can monitor the stack creation in the AWS CloudFormation console.")
 
     except Exception as e:
         logger.error(f"Failed to create cluster stack: {e}")
         if debug:
             logger.exception("Detailed error information:")
         raise click.ClickException(str(e))
 
+
 @click.command("cluster-stack")
 @click.argument("stack-name", required=True)
 @click.option("--region", help="AWS region")
@@ -223,6 +183,7 @@ def describe_cluster_stack(stack_name: str, debug: bool, region: str) -> None:
 
         raise click.ClickException(str(e))
 
+
 @click.command("cluster-stack")
 @click.option("--region", help="AWS region")
 @click.option("--debug", is_flag=True, help="Enable debug logging")
@@ -294,10 +255,11 @@ def list_cluster_stacks(region, debug, status):
 
         raise click.ClickException(str(e))
 
+
 @click.command("cluster-stack")
 @click.argument("stack-name", required=True)
 @click.option("--retain-resources", help="Comma-separated list of logical resource IDs to retain during deletion (only works on DELETE_FAILED stacks). Resource names are shown in failed deletion output, or use AWS CLI: 'aws cloudformation list-stack-resources --stack-name STACK_NAME --region REGION'")
-@click.option("--region", required=True, help="AWS region (required)")
+@click.option("--region", required=True, help="AWS region")
 @click.option("--debug", is_flag=True, help="Enable debug logging")
 @_hyperpod_telemetry_emitter(Feature.HYPERPOD_CLI, "delete_cluster_stack_cli")
 def delete_cluster_stack(stack_name: str, retain_resources: str, region: str, debug: bool) -> None:
@@ -314,6 +276,10 @@ def delete_cluster_stack(stack_name: str, retain_resources: str, region: str, de
           # Delete a cluster stack
           hyp delete cluster-stack my-stack-name --region us-west-2
 
+          # Delete with retained resources (only works on DELETE_FAILED stacks)
+          hyp delete cluster-stack my-stack-name --retain-resources S3Bucket-TrainingData,EFSFileSystem-Models --region us-west-2
+          hyp delete cluster-stack my-stack-name --region us-west-2
+
           # Delete with retained resources (only works on DELETE_FAILED stacks)
           hyp delete cluster-stack my-stack-name --retain-resources S3Bucket-TrainingData,EFSFileSystem-Models --region us-west-2
     """
@@ -329,7 +295,7 @@ def delete_cluster_stack(stack_name: str, retain_resources: str, region: str, de
             confirm_callback=lambda msg: click.confirm("Continue?", default=False),
             success_callback=lambda msg: click.echo(f"✓ {msg}")
         )
-        
+
     except StackNotFoundError:
         click.secho(f"❌ Stack '{stack_name}' not found", fg='red')
     except click.ClickException:
@@ -341,6 +307,7 @@ def delete_cluster_stack(stack_name: str, retain_resources: str, region: str, de
             logger.exception("Detailed error information:")
         raise click.ClickException(str(e))
 
+
 @click.command("cluster")
 @click.option("--cluster-name", required=True, help="The name of the cluster to update")
 @click.option("--instance-groups", help="Instance Groups JSON string")

src/sagemaker/hyperpod/cli/commands/init.py

@@ -8,11 +8,8 @@
 from sagemaker.hyperpod.cli.constants.init_constants import (
     USAGE_GUIDE_TEXT_CFN,
     USAGE_GUIDE_TEXT_CRD,
-    CFN,
-    CRD
+    CFN
 )
-from sagemaker.hyperpod.common.config import Metadata
-from sagemaker.hyperpod.training.hyperpod_pytorch_job import HyperPodPytorchJob
 from sagemaker.hyperpod.cluster_management.hp_cluster_stack import HpClusterStack
 from sagemaker.hyperpod.cli.init_utils import (
     generate_click_command,
@@ -25,8 +22,7 @@
     display_validation_results,
     build_config_from_schema,
     save_template,
-    get_default_version_for_template,
-    add_default_az_ids_to_config,
+    get_default_version_for_template
 )
 from sagemaker.hyperpod.common.utils import get_aws_default_region
 
@@ -265,7 +261,7 @@ def validate():
 
 
 @click.command(name="_default_create")
-@click.option("--region", "-r", default=None, help="Region, default to your region in aws configure")
+@click.option("--region", "-r", default=None, help="Region to create cluster stack for, default to your region in aws configure. Not available for other templates.")
 def _default_create(region):
     """
     Validate configuration and render template files for deployment.
@@ -300,6 +296,11 @@ def _default_create(region):
     # 1) Load config to determine template type
     data, template, version = load_config_and_validate(dir_path)
 
+    # Check if region flag is used for non-cluster-stack templates
+    if region and template != "cluster-stack":
+        click.secho(f"❌  --region flag is only available for cluster-stack template, not for {template}.", fg="red")
+        sys.exit(1)
+    
     # 2) Determine correct jinja file based on template type
     info = TEMPLATES[template]
     schema_type = info["schema_type"]
@@ -327,27 +328,7 @@ def _default_create(region):
     try:
         template_source = jinja_file.read_text()
         tpl = Template(template_source)
-        
-        # For CFN templates, prepare arrays for Jinja template
-        if schema_type == CFN:
-            # Prepare instance_group_settings array
-            instance_group_settings = []
-            rig_settings = []
-            for i in range(1, 21):
-                ig_key = f'instance_group_settings{i}'
-                rig_key = f'rig_settings{i}'
-                if ig_key in data:
-                    instance_group_settings.append(data[ig_key])
-                if rig_key in data:
-                    rig_settings.append(data[rig_key])
-            
-            # Add arrays to template context
-            template_data = dict(data)
-            template_data['instance_group_settings'] = instance_group_settings
-            template_data['rig_settings'] = rig_settings
-            rendered = tpl.render(**template_data)
-        else:
-            rendered = tpl.render(**data)
+        rendered = tpl.render(**data)
     except Exception as e:
         click.secho(f"❌  Failed to render template: {e}", fg="red")
         sys.exit(1)
@@ -375,27 +356,26 @@ def _default_create(region):
             region = get_aws_default_region()
             click.secho(f"Submitting to default region: {region}.", fg="yellow")
 
-        if schema_type == CFN:
-            add_default_az_ids_to_config(out_dir, region)
-
-            from sagemaker.hyperpod.cli.commands.cluster_stack import create_cluster_stack_helper
-            create_cluster_stack_helper(config_file=f"{out_dir}/config.yaml",
-                                        region=region)
-        else:
-            dir_path = Path(".").resolve()
-            data, template, version = load_config(dir_path)
-            namespace = data.get("namespace", "default")
-            registry = TEMPLATES[template]["registry"]
-            model = registry.get(str(version))
-            if model:
-                # Filter out CLI metadata fields before passing to model
-                from sagemaker.hyperpod.cli.init_utils import _filter_cli_metadata_fields
-                filtered_config = _filter_cli_metadata_fields(data)
-                flat = model(**filtered_config)
+        # Unified pattern for all templates
+        dir_path = Path(".").resolve()
+        data, template, version = load_config(dir_path)
+        registry = TEMPLATES[template]["registry"]
+        model = registry.get(str(version))
+        if model:
+            # Filter out CLI metadata fields before passing to model
+            from sagemaker.hyperpod.cli.init_utils import _filter_cli_metadata_fields
+            filtered_config = _filter_cli_metadata_fields(data)
+            flat = model(**filtered_config)
+            
+            # Pass region to to_domain for cluster stack template
+            if template == "cluster-stack":
+                config = flat.to_config(region=region)
+                HpClusterStack(**config).create(region)
+            else:
                 domain = flat.to_domain()
                 domain.create()
 
 
     except Exception as e:
         click.secho(f"❌  Failed to submit the command: {e}", fg="red")
-        sys.exit(1)
+        sys.exit(1)