Update jinja template handling logic for inference and training (#279)

* Update jinja template handling logic for inference and training, cluster logic remaining for discussion

* test inference and training all parameters

* minor change to fix integ

* move create_from_k8s_yaml to init_utils and init_constants, reuse create for create_from_dict for inference SDK, update template kind

* Fix unit test

* update create and create_from_dict for inference to share internal_create, revert unit test changes fix

Author: mollyheamazon

hyperpod-custom-inference-template/hyperpod_custom_inference_template/v1_0/model.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 from pydantic import BaseModel, Field, model_validator, ConfigDict
 from typing import Optional, List, Dict, Union, Literal
+import yaml
 
 from sagemaker.hyperpod.inference.config.hp_endpoint_config import (
     Metrics,
@@ -367,4 +368,4 @@ def to_domain(self) -> HPEndpoint:
             worker=worker,
             invocation_endpoint=self.invocation_endpoint,
             auto_scaling_spec=auto_scaling_spec
-        )
+        )

hyperpod-custom-inference-template/hyperpod_custom_inference_template/v1_0/template.py

@@ -1,68 +1,88 @@
-TEMPLATE_CONTENT = """### Please keep template file unchanged ###
-apiVersion: hyperpod.sagemaker.aws/v1
-kind: HPEndpoint
+TEMPLATE_CONTENT = """
+apiVersion: inference.sagemaker.aws.amazon.com/v1alpha1
+kind: InferenceEndpointConfig
 metadata:
-  name: "{{ endpoint_name }}"
-  namespace: "{{ namespace }}"
+  name: {{ metadata_name or endpoint_name }}
+  namespace: {{ namespace }}
 spec:
-  instanceType: "{{ instance_type }}"
-  modelName:    "{{ model_name }}"
-{% if model_version is not none %}  modelVersion: "{{ model_version }}"
-{% endif %}
-  env:
-{% if env %}    
-{% for key, val in env.items() %}    - name:  "{{ key }}"
-      value: "{{ val }}"
-{% endfor %}{% else %}    []
-{% endif %}
+  endpointName: {{ endpoint_name }}
+  instanceType: {{ instance_type }}
+  modelName: {{ model_name }}
+  modelVersion: {{ model_version or "" }}
+  
   metrics:
-    enabled: {{ metrics_enabled }}
+    enabled: {{ metrics_enabled or False }}
+  
   modelSourceConfig:
-    modelSourceType: "{{ model_source_type }}"
-{% if model_location is not none %}    modelLocation:   "{{ model_location }}"
-{% endif %}    prefetchEnabled: {{ prefetch_enabled }}
-{% if model_source_type == "s3" %}    s3Storage:
-      bucketName: "{{ s3_bucket_name }}"
-      region:     "{{ s3_region }}"
-{% elif model_source_type == "fsx" %}    fsxStorage:
-      dnsName:       "{{ fsx_dns_name }}"
-      fileSystemId:  "{{ fsx_file_system_id }}"
-{% if fsx_mount_name is not none %}      mountName:     "{{ fsx_mount_name }}"
-{% endif %}{% endif %}
+    modelSourceType: {{ model_source_type }}
+    modelLocation: {{ model_location or "" }}
+    prefetchEnabled: {{ prefetch_enabled or False }}
+{%- if model_source_type == "s3" %}
+    s3Storage:
+      bucketName: {{ s3_bucket_name }}
+      region: {{ s3_region }}
+{%- elif model_source_type == "fsx" %}
+    fsxStorage:
+      dnsName: {{ fsx_dns_name }}
+      fileSystemId: {{ fsx_file_system_id }}
+      mountName: {{ fsx_mount_name or "" }}
+{%- endif %}
+  
   tlsConfig:
-{% if tls_certificate_output_s3_uri is not none %}    certificateOutputS3Uri: "{{ tls_certificate_output_s3_uri }}"
-{% else %}    {}
-{% endif %}
+    tlsCertificateOutputS3Uri: {{ tls_certificate_output_s3_uri or "" }}
+  
   worker:
-    image:         "{{ image_uri }}"
-    containerPort: {{ container_port }}
-    volumeMount:
-      name:       "{{ model_volume_mount_name }}"
-      mountPath:  "{{ model_volume_mount_path }}"
+    environmentVariables:
+  {%- if env %}
+  {%- for key, val in env.items() %}
+      - name: {{ key }}
+        value: "{{ val }}"
+  {%- endfor %}
+  {%- else %}
+      []
+  {%- endif %}
+    image: {{ image_uri }}
+    modelInvocationPort:
+      containerPort: {{ container_port }}
+    modelVolumeMount:
+      name: {{ model_volume_mount_name }}
+      mountPath: {{ model_volume_mount_path }}
     resources:
-{% if resources_limits %}      limits:
-{% for key, val in resources_limits.items() %}        {{ key }}: "{{ val }}"
-{% endfor %}{% else %}      {}
-{% endif %}{% if resources_requests %}
+{%- if resources_limits %}
+      limits:
+{%-   for key, val in resources_limits.items() %}
+        {{ key }}: {{ val }}
+{%-   endfor %}
+{%- else %}
+      {}
+{%- endif %}
+{%- if resources_requests %}
       requests:
-{% for key, val in resources_requests.items() %}        {{ key }}: "{{ val }}"
-{% endfor %}{% endif %}
+{%-   for key, val in resources_requests.items() %}
+        {{ key }}: {{ val }}
+{%-   endfor %}
+{%- endif %}
+  
   autoScalingSpec:
     cloudWatchTrigger:
-{% if dimensions %}      dimensions:
-{% for dim_key, dim_val in dimensions.items() %}        - name:  "{{ dim_key }}"
-          value: "{{ dim_val }}"
-{% endfor %}{% else %}      []
-{% endif %}      metricCollectionPeriod: {{ metric_collection_period }}
+{%- if dimensions %}
+      dimensions:
+{%-   for dim_key, dim_val in dimensions.items() %}
+        - name: {{ dim_key }}
+          value: {{ dim_val }}
+{%-   endfor %}
+{%- endif %}
+      metricCollectionPeriod: {{ metric_collection_period }}
       metricCollectionStartTime: {{ metric_collection_start_time }}
-      metricName: "{{ metric_name }}"
-      metricStat: "{{ metric_stat }}"
-      type:       "{{ metric_type }}"
-      minValue:   {{ min_value }}
-      name:       "{{ cloud_watch_trigger_name }}"
-      namespace:  "{{ cloud_watch_trigger_namespace }}"
-      targetValue: {{ target_value }}
-      useCachedMetrics: {{ use_cached_metrics }}
-  invocationEndpoint: "{{ invocation_endpoint }}"
+      metricName: {{ metric_name or "" }}
+      metricStat: {{ metric_stat }}
+      metricType: {{ metric_type }}
+      minValue: {{ min_value }}
+      name: {{ cloud_watch_trigger_name or "" }}
+      namespace: {{ cloud_watch_trigger_namespace or "" }}
+      targetValue: {{ target_value or "" }}
+      useCachedMetrics: {{ use_cached_metrics or False }}
+  
+  invocationEndpoint: {{ invocation_endpoint }}
 
 """

hyperpod-jumpstart-inference-template/hyperpod_jumpstart_inference_template/v1_0/model.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 from pydantic import BaseModel, Field, model_validator, ConfigDict
 from typing import Optional
+import yaml
 
 # reuse the nested types
 from sagemaker.hyperpod.inference.config.hp_jumpstart_endpoint_config import (
@@ -114,4 +115,4 @@ def to_domain(self) -> HPJumpStartEndpoint:
             server=server,
             sage_maker_endpoint=sage_ep,
             tls_config=tls
-        )
+        )

hyperpod-jumpstart-inference-template/hyperpod_jumpstart_inference_template/v1_0/template.py

@@ -1,8 +1,8 @@
-TEMPLATE_CONTENT = """### Please keep template file unchanged ###
+TEMPLATE_CONTENT = """
 apiVersion: inference.sagemaker.aws.amazon.com/v1alpha1
 kind: JumpStartModel
 metadata:
-  name:                {{ model_id }}
+  name:                {{ metadata_name or endpoint_name }}
   namespace:           {{ namespace or "default" }}
 spec:
   model:
@@ -14,4 +14,6 @@
     name:                     {{ endpoint_name or "" }}
   server:
     instanceType:             {{ instance_type }}
+  tlsConfig:
+    tlsCertificateOutputS3Uri: {{ tls_certificate_output_s3_uri or "" }}
 """

hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/v1_0/model.py

@@ -14,7 +14,7 @@
     PersistentVolumeClaim
 )
 from sagemaker.hyperpod.training.hyperpod_pytorch_job import HyperPodPytorchJob
-
+import yaml
 
 class VolumeConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -39,7 +39,7 @@ class VolumeConfig(BaseModel):
         description="PVC claim name (required for pvc volumes)",
         min_length=1
     )
-    read_only: Optional[Literal['true', 'false']] = Field(None, description="Read-only flag for pvc volumes")
+    read_only: Optional[bool] = Field(None, description="Read-only flag for pvc volumes")
 
     @field_validator('mount_path', 'path')
     @classmethod
@@ -260,7 +260,7 @@ def build_dict(**kwargs):
                 elif vol.type == "pvc":
                     volume_obj = Volumes(name=vol.name, persistent_volume_claim=PersistentVolumeClaim(
                         claim_name=vol.claim_name,
-                        read_only=vol.read_only == "true" if vol.read_only else False
+                        read_only=vol.read_only if vol.read_only is not None else False
                     ))
                 volumes.append(volume_obj)
 
@@ -310,6 +310,7 @@ def build_dict(**kwargs):
         result = HyperPodPytorchJob(**job_kwargs)
         return result
 
+
 # Volume-specific type handlers - only override what's needed
 def volume_parse_strings(ctx_or_strings, param=None, value=None):
     """Parse volume strings into VolumeConfig objects. Can be used as Click callback."""

hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/v1_0/schema.json

@@ -54,11 +54,7 @@
         "read_only": {
           "anyOf": [
             {
-              "enum": [
-                "true",
-                "false"
-              ],
-              "type": "string"
+              "type": "boolean"
             },
             {
               "type": "null"

hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/v1_0/template.py

@@ -1,68 +1,96 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License"). You
-# may not use this file except in compliance with the License. A copy of
-# the License is located at
-#
-#     http://aws.amazon.com/apache2.0/
-#
-# or in the "license" file accompanying this file. This file is
-# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
-# ANY KIND, either express or implied. See the License for the specific
-# language governing permissions and limitations under the License.
-TEMPLATE_CONTENT = """### Please keep template file unchanged ###
+TEMPLATE_CONTENT = """
 apiVersion: sagemaker.amazonaws.com/v1
 kind: HyperPodPyTorchJob
 metadata:
-  name: "{{ job_name }}"
-  namespace: "{{ namespace }}"
-{% if queue_name or priority %}  labels:
-{% if queue_name %}    kueue.x-k8s.io/queue-name: "{{ queue_name }}"
-{% endif %}{% if priority %}    kueue.x-k8s.io/priority-class: "{{ priority }}"
-{% endif %}{% endif %}spec:
-{% if tasks_per_node %}  nprocPerNode: "{{ tasks_per_node }}"
-{% endif %}  replicaSpecs:
-    - name: "pod"
-{% if node_count %}      replicas: {{ node_count }}
-{% endif %}      template:
+  name: {{ job_name }}
+  namespace: {{ namespace }}
+{%- if queue_name or priority %}
+  labels:
+    kueue.x-k8s.io/queue-name: {{ queue_name or "" }}
+    kueue.x-k8s.io/priority-class: {{ priority or "" }}
+{%- endif %}
+spec:
+{%- if tasks_per_node %}
+  nprocPerNode: "{{ tasks_per_node }}"
+{%- endif %}
+  replicaSpecs:
+    - name: pod
+      replicas: {{ node_count or 1 }}
+      template:
         metadata:
-          name: "{{ job_name }}"
-{% if namespace %}          namespace: "{{ namespace }}"
-{% endif %}{% if queue_name or priority %}          labels:
-{% if queue_name %}            kueue.x-k8s.io/queue-name: "{{ queue_name }}"
-{% endif %}{% if priority %}            kueue.x-k8s.io/priority-class: "{{ priority }}"
-{% endif %}{% endif %}        spec:
+          name: {{ job_name }}
+          namespace: {{ namespace }}
+{%-       if queue_name or priority %}
+          labels:
+            kueue.x-k8s.io/queue-name: {{ queue_name or "" }}
+            kueue.x-k8s.io/priority-class: {{ priority or "" }}
+{%-       endif %}
+        spec:
           containers:
-            - name: "container-name"
-              image: "{{ image }}"
-{% if pull_policy %}              imagePullPolicy: "{{ pull_policy }}"
-{% endif %}{% if command %}              command: {{ command | tojson }}
-{% endif %}{% if args %}              args: {{ args | tojson }}
-{% endif %}{% if environment %}              env:
-{% for key, value in environment.items() %}                - name: "{{ key }}"
+            - name: container-name
+              image: {{ image }}
+{%-           if pull_policy %}
+              imagePullPolicy: {{ pull_policy }}
+{%-           endif %}
+{%-           if command %}
+              command: {{ command | tojson }}
+{%-           endif %}
+{%-           if args %}
+              args: {{ args | tojson }}
+{%-           endif %}
+{%-           if environment %}
+              env:
+{%-             for key, value in environment.items() %}
+                - name: {{ key }}
                   value: "{{ value }}"
-{% endfor %}{% endif %}{% if volume %}              volumeMounts:
-{% for vol in volume %}                - name: "{{ vol.name }}"
-                  mountPath: "{{ vol.mount_path }}"
-{% if vol.read_only is not none and vol.read_only != "" %}                  readOnly: {{ vol.read_only | lower }}
-{% endif %}{% endfor %}{% endif %}              resources:
+{%-             endfor %}
+{%-           endif %}
+{%-           if volume %}
+              volumeMounts:
+{%-             for vol in volume %}
+                - name: {{ vol.name }}
+                  mountPath: {{ vol.mount_path }}
+                  readOnly: {{ vol.read_only | lower if vol.read_only else false }}
+{%-             endfor %}
+{%-           endif %}
+              resources:
                 requests:
                   nvidia.com/gpu: "0"
                 limits:
                   nvidia.com/gpu: "0"
-{% if instance_type or label_selector or deep_health_check_passed_nodes_only %}          nodeSelector:
-{% if instance_type %}            node.kubernetes.io/instance-type: "{{ instance_type }}"
-{% endif %}{% if label_selector %}{% for key, value in label_selector.items() %}            {{ key }}: "{{ value }}"
-{% endfor %}{% endif %}{% if deep_health_check_passed_nodes_only %}            deep-health-check-passed: "true"
-{% endif %}{% endif %}{% if service_account_name %}          serviceAccountName: "{{ service_account_name }}"
-{% endif %}{% if scheduler_type %}          schedulerName: "{{ scheduler_type }}"
-{% endif %}{% if volume %}          volumes:
-{% for vol in volume %}            - name: "{{ vol.name }}"
-{% if vol.type == "hostPath" %}              hostPath:
-                path: "{{ vol.path }}"
-{% elif vol.type == "pvc" %}              persistentVolumeClaim:
-                claimName: "{{ vol.claim_name }}"
-{% endif %}{% endfor %}{% endif %}{% if max_retry %}  runPolicy:
+{%-         if instance_type or label_selector or deep_health_check_passed_nodes_only %}
+          nodeSelector:
+            node.kubernetes.io/instance-type: {{ instance_type or "" }}
+{%-           if label_selector %}
+{%-             for key, value in label_selector.items() %}
+            {{ key }}: {{ value }}
+{%-             endfor %}
+{%-           endif %}
+{%-           if deep_health_check_passed_nodes_only %}
+            deep-health-check-passed: "true"
+{%-           endif %}
+{%-         endif %}
+{%-         if service_account_name %}
+          serviceAccountName: {{ service_account_name }}
+{%-         endif %}
+{%-         if scheduler_type %}
+          schedulerName: {{ scheduler_type }}
+{%-         endif %}
+{%-         if volume %}
+          volumes:
+{%-           for vol in volume %}
+            - name: {{ vol.name }}
+{%-             if vol.type == "hostPath" %}
+              hostPath:
+                path: {{ vol.path }}
+{%-             elif vol.type == "pvc" %}
+              persistentVolumeClaim:
+                claimName: {{ vol.claim_name }}
+{%-             endif %}
+{%-           endfor %}
+{%-         endif %}
+{%- if max_retry %}
+  runPolicy:
     cleanPodPolicy: "None"
     jobMaxRetryCount: {{ max_retry }}
-{% endif %}"""
+{%- endif %}"""