Modify efa arg name and fix gpu integ test

Author: yungwenh-aws

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

@@ -195,12 +195,12 @@ class PyTorchJobConfig(BaseModel):
         default=None,
         description="Limit for the amount of memory in GiB",
     )
-    efa: Optional[int] = Field(
+    efa_interfaces: Optional[int] = Field(
         default=None,
         description="Number of EFA interfaces for the instance",
         ge=0
     )
-    efa_limit: Optional[int] = Field(
+    efa_interfaces_limit: Optional[int] = Field(
         default=None,
         description="Limit for the number of EFA interfaces",
         ge=0
@@ -464,26 +464,26 @@ def build_dict(**kwargs):
                     **{partition_resource_key: str(self.accelerator_partition_count)} if self.accelerator_partition_count else {},
                     vcpu=str(self.vcpu) if self.vcpu else None,
                     memory=str(self.memory) if self.memory else None,
-                    **{"vpc.amazonaws.com/efa": str(self.efa)} if self.efa else {},
+                    **{"vpc.amazonaws.com/efa": str(self.efa_interfaces)} if self.efa_interfaces else {},
                 )
                 limits_value = build_dict(
                     **{partition_resource_key: str(self.accelerator_partition_limit)} if self.accelerator_partition_limit else {},
                     vcpu=str(self.vcpu_limit) if self.vcpu_limit else None,
                     memory=str(self.memory_limit) if self.memory_limit else None,
-                    **{"vpc.amazonaws.com/efa": str(self.efa_limit)} if self.efa_limit else {},
+                    **{"vpc.amazonaws.com/efa": str(self.efa_interfaces_limit)} if self.efa_interfaces_limit else {},
                 )
             else:
                 requests_value = build_dict(
                     accelerators=str(self.accelerators) if self.accelerators else None,
                     vcpu=str(self.vcpu) if self.vcpu else None,
                     memory=str(self.memory) if self.memory else None,
-                    **{"vpc.amazonaws.com/efa": str(self.efa)} if self.efa else {},
+                    **{"vpc.amazonaws.com/efa": str(self.efa_interfaces)} if self.efa_interfaces else {},
                 )
                 limits_value = build_dict(
                     accelerators=str(self.accelerators_limit) if self.accelerators_limit else None,
                     vcpu=str(self.vcpu_limit) if self.vcpu_limit else None,
                     memory=str(self.memory_limit) if self.memory_limit else None,
-                    **{"vpc.amazonaws.com/efa": str(self.efa_limit)} if self.efa_limit else {},
+                    **{"vpc.amazonaws.com/efa": str(self.efa_interfaces_limit)} if self.efa_interfaces_limit else {},
                 )
 
         # Build container

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

@@ -305,12 +305,12 @@
       "minimum": 0,
       "description": "Limit for the amount of memory in GiB"
     },
-    "efa": {
+    "efa_interfaces": {
       "type": "integer",
       "minimum": 0,
       "description": "Number of EFA interfaces for the instance"
     },
-    "efa_limit": {
+    "efa_interfaces_limit": {
       "type": "integer",
       "minimum": 0,
       "description": "Limit for the number of EFA interfaces"

src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py

@@ -152,15 +152,15 @@ def _process_replica_resources(cls, data):
             acc_req, acc_lim = _set_default_accelerators_val(instance_type, accelerators, accelerators_limit)
             _validate_accelerators_inputs(instance_type, acc_req, acc_lim)
 
-            efa = None
+            efa_interfaces = None
             if requests.get(EFA_RESOURCE_KEY):
-                efa = int(requests.get(EFA_RESOURCE_KEY))
+                efa_interfaces = int(requests.get(EFA_RESOURCE_KEY))
 
-            efa_limit = None
+            efa_interfaces_limit = None
             if limits.get(EFA_RESOURCE_KEY):
-                efa_limit = int(limits.get(EFA_RESOURCE_KEY))
+                efa_interfaces_limit = int(limits.get(EFA_RESOURCE_KEY))
 
-            _validate_efa_inputs(instance_type, efa, efa_limit)
+            _validate_efa_inputs(instance_type, efa_interfaces, efa_interfaces_limit)
 
             accelerator_partition_type, accelerator_partition_count, accelerator_partition_limit = (
                 _get_accelerator_partition(requests, limits)
@@ -174,7 +174,7 @@ def _process_replica_resources(cls, data):
 
             acc_partition_req, acc_partition_lim = _set_default_accelerator_partition_val(accelerator_partition_count, accelerator_partition_limit)
 
-            requests_values = _get_resources_from_compute_quotas(instance_type, vcpu, memory, acc_req, accelerator_partition_type, acc_partition_req, efa)
+            requests_values = _get_resources_from_compute_quotas(instance_type, vcpu, memory, acc_req, accelerator_partition_type, acc_partition_req, efa_interfaces)
             if requests_values is None:
                 requests_values = _get_resources_from_instance(instance_type, node_count=1)
                 _trim_resource_requests(instance_type, requests_values)

src/sagemaker/hyperpod/training/quota_allocation_util.py

@@ -34,7 +34,7 @@ def _get_resources_from_compute_quotas(instance_type: str,
                                        accelerators: Optional[int] = 0,
                                        accelerator_partition_type: Optional[str] = None,
                                        accelerator_partition_count: Optional[int] = None,
-                                       efa: Optional[int] = None) -> Optional[dict]:
+                                       efa_interfaces: Optional[int] = None) -> Optional[dict]:
     has_accelerator_partition = accelerator_partition_type is not None and accelerator_partition_count is not None
     has_compute_resources = _has_compute_resource_quota_allocation_resources(memory_in_gib, vcpu, accelerators)
 
@@ -73,7 +73,7 @@ def _get_resources_from_compute_quotas(instance_type: str,
         result["memory"] = memory_value
         result[type_of_accelerator] = accelerators
 
-        efa_count = efa or instance.get("efa", 0)
+        efa_count = efa_interfaces or instance.get("efa", 0)
         if efa_count > 0:
             result["vpc.amazonaws.com/efa"] = efa_count
 
@@ -135,7 +135,7 @@ def _trim_resource_requests(instance_type: str, requests_values: dict) -> dict:
     return requests_values
 
 
-def _get_limits(instance_type: str, vcpu_limit: Optional[float], memory_in_gib_limit: Optional[float], accelerators_limit: Optional[int], accelerator_partition_type: Optional[str], accelerator_partition_limit: Optional[int], efa_limit: Optional[int] = None) -> dict:
+def _get_limits(instance_type: str, vcpu_limit: Optional[float], memory_in_gib_limit: Optional[float], accelerators_limit: Optional[int], accelerator_partition_type: Optional[str], accelerator_partition_limit: Optional[int], efa_interfaces_limit: Optional[int] = None) -> dict:
 
     result = {}
     type_of_accelerator, _max_accelerator_per_instance = _get_accelerator_type_and_count(instance_type)
@@ -154,8 +154,8 @@ def _get_limits(instance_type: str, vcpu_limit: Optional[float], memory_in_gib_l
     if memory_in_gib_limit is not None:
         result["memory"] = str(memory_in_gib_limit) + "Gi"
 
-    if efa_limit is not None and efa_limit > 0:
-        result["vpc.amazonaws.com/efa"] = efa_limit
+    if efa_interfaces_limit is not None and efa_interfaces_limit > 0:
+        result["vpc.amazonaws.com/efa"] = efa_interfaces_limit
 
     return result
 
@@ -226,29 +226,29 @@ def _validate_accelerators_inputs(instance_type: str, accelerators_request: int,
                 raise ValueError('Requested accelerators exceeds capacity')
 
 
-def _validate_efa_inputs(instance_type: str, efa_request: Optional[int], efa_limit: Optional[int]) -> None:
+def _validate_efa_inputs(instance_type: str, efa_interfaces: Optional[int], efa_interfaces_limit: Optional[int]) -> None:
     """Validate EFA inputs similar to accelerator validation."""
     instance = INSTANCE_RESOURCES.get(instance_type, {})
     max_efa_per_instance = instance.get("efa", 0)
 
     # Check if user provided EFA values but instance doesn't support EFA
-    if max_efa_per_instance == 0 and (efa_request is not None or efa_limit is not None):
+    if max_efa_per_instance == 0 and (efa_interfaces is not None or efa_interfaces_limit is not None):
         raise ValueError(
             f"Instance type {instance_type} does not support EFA, but EFA values were provided.")
 
     # Validate EFA values if instance supports EFA
     if max_efa_per_instance > 0:
-        if efa_request is not None and efa_limit is not None:
-            if efa_request != efa_limit:
+        if efa_interfaces is not None and efa_interfaces_limit is not None:
+            if efa_interfaces != efa_interfaces_limit:
                 raise ValueError('EFA request must equal EFA limit')
-            if efa_limit > max_efa_per_instance:
-                raise ValueError(f'Requested EFA limit ({efa_limit}) exceeds instance capacity ({max_efa_per_instance})')
-            if efa_request > max_efa_per_instance:
-                raise ValueError(f'Requested EFA ({efa_request}) exceeds instance capacity ({max_efa_per_instance})')
-        elif efa_request is not None and efa_request > max_efa_per_instance:
-            raise ValueError(f'Requested EFA ({efa_request}) exceeds instance capacity ({max_efa_per_instance})')
-        elif efa_limit is not None and efa_limit > max_efa_per_instance:
-            raise ValueError(f'Requested EFA limit ({efa_limit}) exceeds instance capacity ({max_efa_per_instance})')
+            if efa_interfaces_limit > max_efa_per_instance:
+                raise ValueError(f'Requested EFA limit ({efa_interfaces_limit}) exceeds instance capacity ({max_efa_per_instance})')
+            if efa_interfaces > max_efa_per_instance:
+                raise ValueError(f'Requested EFA ({efa_interfaces}) exceeds instance capacity ({max_efa_per_instance})')
+        elif efa_interfaces is not None and efa_interfaces > max_efa_per_instance:
+            raise ValueError(f'Requested EFA ({efa_interfaces}) exceeds instance capacity ({max_efa_per_instance})')
+        elif efa_interfaces_limit is not None and efa_interfaces_limit > max_efa_per_instance:
+            raise ValueError(f'Requested EFA limit ({efa_interfaces_limit}) exceeds instance capacity ({max_efa_per_instance})')
 
 
 def _set_default_accelerators_val(instance_type: Optional[str], accelerators_request: Optional[int], accelerators_limit: Optional[int]) -> Tuple[Optional[int], Optional[int]]:

test/integration_tests/training/cli/test_gpu_quota_allocation.py

@@ -53,8 +53,8 @@ def test_create_job_with_integer_quota_parameters(self, test_job_name):
         result = execute_command(describe_cmd)
         logger.info(f"describe result: {result}")
         assert result.returncode == 0
-        assert "      Limits:   {'cpu': '4', 'memory': '2Gi', 'nvidia.com/gpu': '1'}" in result.stdout
-        assert "      Requests: {'cpu': '3', 'memory': '1Gi', 'nvidia.com/gpu': '1'}" in result.stdout
+        assert "      Limits:   {'cpu': '4', 'memory': '2Gi', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
+        assert "      Requests: {'cpu': '3', 'memory': '1Gi', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
 
         delete_cmd = [
             "hyp", "delete", "hyp-pytorch-job",
@@ -103,8 +103,8 @@ def test_create_job_with_float_quota_parameters(self, test_job_name):
         ]
         result = execute_command(describe_cmd)
         assert result.returncode == 0
-        assert "      Limits:   {'cpu': '4800m', 'memory': '2899102924800m', 'nvidia.com/gpu': '1'}" in result.stdout
-        assert "      Requests: {'cpu': '3600m', 'memory': '1Gi', 'nvidia.com/gpu': '1'}" in result.stdout
+        assert "      Limits:   {'cpu': '4800m', 'memory': '2899102924800m', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
+        assert "      Requests: {'cpu': '3600m', 'memory': '1Gi', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
 
         delete_cmd = [
             "hyp", "delete", "hyp-pytorch-job",
@@ -149,8 +149,8 @@ def test_create_job_with_only_accelerators_parameter(self, test_job_name):
         ]
         result = execute_command(describe_cmd)
         assert result.returncode == 0
-        assert "      Limits:   {'memory': '104Gi', 'nvidia.com/gpu': '1'}" in result.stdout
-        assert "      Requests: {'cpu': '29', 'memory': '104Gi', 'nvidia.com/gpu': '1'}" in result.stdout
+        assert "      Limits:   {'memory': '104Gi', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
+        assert "      Requests: {'cpu': '29', 'memory': '104Gi', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
 
         delete_cmd = [
             "hyp", "delete", "hyp-pytorch-job",
@@ -196,8 +196,8 @@ def test_create_job_with_accelerators_memory_parameters(self, test_job_name):
         time.sleep(5)
 
         assert result.returncode == 0
-        assert "      Limits:   {'memory': '2899102924800m', 'nvidia.com/gpu': '1'}" in result.stdout
-        assert "      Requests: {'cpu': '29', 'memory': '2040109465600m', 'nvidia.com/gpu': '1'}" in result.stdout
+        assert "      Limits:   {'memory': '2899102924800m', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
+        assert "      Requests: {'cpu': '29', 'memory': '2040109465600m', 'nvidia.com/gpu': '1', 'vpc.amazonaws.com/efa': '1'}" in result.stdout
 
         delete_cmd = [
             "hyp", "delete", "hyp-pytorch-job",

test/unit_tests/training/test_pytorch_job_template_model.py

@@ -97,7 +97,7 @@ def test_user_specified_efa_overrides_default(self):
             job_name="test-custom-efa",
             image="pytorch:latest",
             accelerators=4,
-            efa=2,
+            efa_interfaces=2,
             instance_type="ml.p4d.24xlarge"
         )