Implementing Task Gov. feature for SDK flow

Author: jam-jee

hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/v1_1/quota_allocation_util.py renamed to hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/quota_allocation_util.py

@@ -20,7 +20,7 @@
     'topology.k8s.aws/network-node-layer-2',
     'topology.k8s.aws/network-node-layer-3'
 }
-from .quota_allocation_util import _is_valid, _get_resources_from_compute_quotas, _get_resources_from_instance, _get_limits
+from hyperpod_pytorch_job_template.quota_allocation_util import _is_valid, _get_resources_from_compute_quotas, _get_resources_from_instance, _get_limits
 
 class VolumeConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -111,7 +111,7 @@ class PyTorchJobConfig(BaseModel):
         min_length=1
     )
     node_count: Optional[int] = Field(
-        default=1, 
+        default=None,
         alias="node_count", 
         description="Number of nodes",
         ge=1
@@ -290,7 +290,7 @@ def to_domain(self) -> Dict:
         valid, error = _is_valid(
            self.vcpu, self.memory, self.accelerators, self.node_count, self.instance_type
         )
-        
+
         if not valid:
             raise ValueError(error)
 

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

@@ -202,7 +202,7 @@
           "type": "null"
         }
       ],
-      "default": 1,
+      "default": null,
       "description": "Number of nodes",
       "title": "Node Count"
     },

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

@@ -44,6 +44,7 @@
 SAGEMAKER_MANAGED_CLUSTER_QUEUE_SUFFIX = "-clusterqueue"
 SAGEMAKER_TRAINING_LAUNCHER_DIR = str(Path(__file__).parent.parent / "sagemaker_hyperpod_recipes")
 NVIDIA_GPU_RESOURCE_LIMIT_KEY = "nvidia.com/gpu"
+NEURON_RESOURCE_LIMIT_KEY = "aws.amazon.com/neurondevice"
 AVAILABLE_ACCELERATOR_DEVICES_KEY = "AvailableAcceleratorDevices"
 TOTAL_ACCELERATOR_DEVICES_KEY = "TotalAcceleratorDevices"
 USER_NAME_LABEL_KEY = "sagemaker.user/created-by"

src/sagemaker/hyperpod/cli/constants/command_constants.py

@@ -2979,7 +2979,7 @@ class ReplicaSpec(BaseModel):
 
     name: str = Field(description="The name for the replica set")
     replicas: Optional[int] = Field(
-        default=1,
+        default=0,
         description="Replicas is the desired number of replicas of the given template.",
     )
     spares: Optional[int] = Field(

src/sagemaker/hyperpod/training/config/hyperpod_pytorch_job_unified_config.py

@@ -1,4 +1,7 @@
 from pydantic import ConfigDict, Field
+
+from sagemaker.hyperpod.cli.constants.command_constants import INSTANCE_TYPE_LABEL, NVIDIA_GPU_RESOURCE_LIMIT_KEY, \
+    NEURON_RESOURCE_LIMIT_KEY
 from sagemaker.hyperpod.training.config.hyperpod_pytorch_job_unified_config import (
     _HyperPodPytorchJob, HyperPodPytorchJobStatus
 )
@@ -18,6 +21,9 @@
 import yaml
 import logging
 
+from hyperpod_pytorch_job_template.quota_allocation_util import _is_valid, _get_resources_from_compute_quotas, _get_resources_from_instance, _get_limits
+
+
 
 TRAINING_GROUP = "sagemaker.amazonaws.com"
 API_VERSION = "v1"
@@ -52,6 +58,109 @@ def verify_kube_config(cls):
 
             # Verify Kubernetes version compatibility
             verify_kubernetes_version_compatibility(cls.get_logger())
+    @classmethod
+    def _extract_numeric_value(cls, value):
+        """Extract numeric value from strings like '1.5Gi' -> 1.5"""
+        if not value:
+            return None
+        import re
+        match = re.match(r'^([0-9]*\.?[0-9]+)', str(value))
+        return float(match.group(1)) if match else None
+
+    @classmethod
+    def sanitize_memory(cls, resource):
+        try :
+            if 'memory' in resource:
+                memory = resource['memory']
+                # Case when quotas have been already initialized in CLI layer
+                # ToDo : Cleanup quota initialization in CLI layer and directly use SDK layer for init.
+                memory.replace('GiGi', 'Gi')
+                resource['memory'] = memory
+            return resource
+        except Exception as e:
+            return resource
+
+
+    @classmethod
+    def _process_replica_resources(cls, data):
+        """Process and validate replica resource configuration."""
+        try:
+            node_count = data.get('replicas', None)
+
+            # Extract nested configuration with validation
+            template = data.get('template', {})
+            spec = template.get('spec', {})
+            node_selector = spec.get('nodeSelector', {})
+            instance_type = node_selector.get(INSTANCE_TYPE_LABEL) if node_selector else None
+
+            if not instance_type:
+                return None
+
+            containers = spec.get('containers', [])
+
+            if not containers:
+                raise ValueError("No containers found in template spec")
+
+            container = containers[0]
+            resources = container.get('resources', {})
+            requests = resources.get('requests', {})
+            limits = resources.get('limits', {})
+
+            # Extract resource values
+            vcpu = float(requests.get('cpu')) if requests.get('cpu') else None
+            memory = cls._extract_numeric_value(requests.get('memory'))
+            accelerators = int(requests.get(NVIDIA_GPU_RESOURCE_LIMIT_KEY)) or int(requests.get(NEURON_RESOURCE_LIMIT_KEY))  or None
+            memory_limit = cls._extract_numeric_value(limits.get('memory'))
+            vcpu_limit = float(limits.get('cpu')) if limits.get('cpu') else None
+            accelerators_limit = int(limits.get(NVIDIA_GPU_RESOURCE_LIMIT_KEY)) or int(limits.get(NEURON_RESOURCE_LIMIT_KEY))  or None
+
+            # Validate configuration
+            valid, error = _is_valid(vcpu, memory, accelerators, node_count, instance_type)
+            if not valid:
+                raise ValueError(error)
+
+            # Calculate resource values
+            requests_value = (_get_resources_from_compute_quotas(instance_type, vcpu, memory, accelerators)
+                              or _get_resources_from_instance(instance_type, node_count))
+            limits_value = _get_limits(instance_type, vcpu_limit, memory_limit, accelerators_limit)
+
+            requests_value = cls.sanitize_memory(requests_value)
+            limits_value = cls.sanitize_memory(limits_value)
+
+            # Update data with calculated values
+            data['template']['spec']['containers'][0]['resources']['requests'] = requests_value
+            data['template']['spec']['containers'][0]['resources']['limits'] = limits_value
+            return data
+        except KeyError as e:
+            raise ValueError(f"Missing required configuration key: {str(e)}")
+
+    @classmethod
+    def _get_container_resources(cls, replica_spec):
+        """Extract container resources from replica spec."""
+        container_resources = replica_spec['template']['spec']['containers'][0]['resources']
+        return container_resources['requests'], container_resources['limits']
+
+    @classmethod
+    def allocate_quotas_if_applicable(cls, spec):
+        logger = cls.get_logger()
+        logger = setup_logging(logger)
+        try:
+            spec_dict = spec.model_dump()
+            replica_spec = spec_dict['replicaSpecs'][0]
+            cls._process_replica_resources(replica_spec)
+
+            # Update the original spec object directly
+            requests, limits = cls._get_container_resources(replica_spec)
+            spec.replicaSpecs[0].template.spec.containers[0].resources.requests = requests
+            spec.replicaSpecs[0].template.spec.containers[0].resources.limits = limits
+
+            return spec
+        except ValueError as e:
+            logger.error(f"Error: in quota allocation:{e}")
+            raise ValueError(e)
+        except Exception as e:
+            logger.info(f"Warning: in quota allocation: {e}. using defaults.")
+            return spec
 
     @_hyperpod_telemetry_emitter(Feature.HYPERPOD, "create_pytorchjob")
     def create(self, debug=False):
@@ -65,6 +174,10 @@ def create(self, debug=False):
         if not self.metadata.namespace:
             self.metadata.namespace = get_default_namespace()
 
+        spec = self.allocate_quotas_if_applicable(spec)
+        if spec.replicaSpecs[0].replicas == 0 :
+            spec.replicaSpecs[0].replicas = 1 # default value
+
         config = {
             "apiVersion": f"{TRAINING_GROUP}/{API_VERSION}",
             "kind": KIND,
@@ -91,6 +204,8 @@ def create(self, debug=False):
             logger.error(f"Failed to create HyperPodPytorchJob {self.metadata.name}!")
             handle_exception(e, self.metadata.name, self.metadata.namespace)
 
+
+
     @classmethod
     @_hyperpod_telemetry_emitter(Feature.HYPERPOD, "list_pytorchjobs")
     def list(cls, namespace=None) -> List["HyperPodPytorchJob"]: