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")

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

@@ -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,88 @@ def verify_kube_config(cls):
 
             # Verify Kubernetes version compatibility
             verify_kubernetes_version_compatibility(cls.get_logger())
+    @classmethod
+    def sanitize_memory(cls, resource):
+        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
+
+    @classmethod
+    def _process_replica_resources(cls, data):
+        """Process and validate replica resource configuration."""
+        try:
+            node_count = data['replicas']
+
+            # Extract nested configuration with validation
+            template = data.get('template', {})
+            spec = template.get('spec', {})
+            node_selector = spec.get('nodeSelector', {})
+            containers = spec.get('containers', [])
+
+            if not containers:
+                raise ValueError("No containers found in template spec")
+
+            instance_type = node_selector.get(INSTANCE_TYPE_LABEL, None)
+            if not instance_type:
+                raise ValueError("Instance type not found in node selector")
+
+            container = containers[0]
+            resources = container.get('resources', {})
+            requests = resources.get('requests', {})
+            limits = resources.get('limits', {})
+
+            # Extract resource values
+            vcpu = requests.get('vcpu', None)
+            memory = requests.get('memory', None)
+            accelerators = requests.get(NVIDIA_GPU_RESOURCE_LIMIT_KEY) or requests.get(NEURON_RESOURCE_LIMIT_KEY) or None
+            memory_limit = limits.get('memory', None)
+            vcpu_limit = limits.get('vcpu', None)
+            accelerators_limit = limits.get(NVIDIA_GPU_RESOURCE_LIMIT_KEY) or requests.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.sanitze_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):
+        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 Exception as e:
+            print(f"Warning: in quota allocation: {e}. using defaults.")
+            return spec
 
     @_hyperpod_telemetry_emitter(Feature.HYPERPOD, "create_pytorchjob")
     def create(self, debug=False):
@@ -65,6 +153,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 +183,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"]: