fix: handle resource units during training quota allocation

doc/getting_started/training.md

@@ -267,6 +267,20 @@ pytorch_job.create()
 ````
 `````
 
+### SDK resource quantities
+
+When `node_selector` includes `node.kubernetes.io/instance-type`, SDK quota
+allocation interprets `Resources.requests` and `Resources.limits` as Kubernetes
+quantities. For example, `{"cpu": "500m", "memory": "1024Mi"}` requests 0.5 CPU
+cores and 1 GiB; limits of `{"cpu": "1", "memory": "2048Mi"}` mean 1 core and
+2 GiB. Unitless memory values retain their legacy GiB meaning, including exponent
+notation, so existing SDK configurations and installed template packages keep
+their allocation behavior. Explicit suffixes such as `Mi`, `G`, and `Ki` use
+Kubernetes units. Invalid, negative, or non-finite quantities raise `ValueError`
+before allocation updates the job.
+
+CLI `--memory` and `--memory-limit` parameters remain amounts in GiB.
+
 ### Key Parameters
 
 When creating a training job, you'll need to specify:

src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py

@@ -1,4 +1,7 @@
 from pydantic import ConfigDict, Field
+from decimal import Decimal, DecimalException, InvalidOperation
+from math import isfinite
+from kubernetes.utils.quantity import parse_quantity
 
 from sagemaker.hyperpod.cli.constants.command_constants import (
     INSTANCE_TYPE_LABEL,
@@ -11,7 +14,7 @@
 )
 from sagemaker.hyperpod.common.config.metadata import Metadata
 from kubernetes import client, config, stream
-from typing import List, Optional, ClassVar
+from typing import List, Optional, ClassVar, Union
 from sagemaker.hyperpod.common.utils import (
     handle_exception,
     get_default_namespace,
@@ -85,17 +88,40 @@ 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"""
+    def _parse_resource_quantity(
+        cls, value: Optional[Union[int, str]], resource_name: str
+    ) -> Optional[float]:
+        """Normalize resource units, retaining legacy unitless memory in GiB.
+
+        Compatible template packages emit unitless GiB values, including
+        exponent notation. Explicit memory suffixes use Kubernetes units.
+        Missing or numeric-zero values retain the existing default behavior.
+        Invalid, negative, or non-finite values raise ValueError before the
+        job is changed.
+        """
         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
+        message = f"Invalid {resource_name} quantity: {value!r}"
+        try:
+            if resource_name == "memory":
+                try:
+                    quantity = Decimal(str(value))
+                except InvalidOperation:
+                    quantity = parse_quantity(value) / (1024**3)
+            else:
+                quantity = parse_quantity(value)
+        except (ValueError, DecimalException) as exc:
+            raise ValueError(message) from exc
+        if not quantity.is_finite() or quantity < 0:
+            raise ValueError(message)
+        normalized = float(quantity)
+        if not isfinite(normalized):
+            raise ValueError(message)
+        return normalized
 
     @classmethod
     def _process_replica_resources(cls, data):
-        """Process and validate replica resource configuration."""
+        """Validate replica resources, interpreting CPU in cores and memory in GiB."""
         try:
             node_count = data.get('replicas', None)
 
@@ -134,20 +160,14 @@ def _process_replica_resources(cls, data):
                 accelerators = int(requests.get(NEURON_RESOURCE_KEY))
 
             # Extract resource values
-            vcpu = None
-            if requests.get('cpu'):
-                vcpu = float(requests.get('cpu'))
-            elif requests.get('vcpu'):
-                vcpu = float(requests.get('vcpu'))
-
-            vcpu_limit = None
-            if limits.get('cpu'):
-                vcpu_limit = float(limits.get('cpu'))
-            elif limits.get('vcpu'):
-                vcpu_limit = float(limits.get('vcpu'))
-
-            memory = cls._extract_numeric_value(requests.get('memory'))
-            memory_limit = cls._extract_numeric_value(limits.get('memory'))
+            vcpu = cls._parse_resource_quantity(
+                requests.get('cpu', requests.get('vcpu')), 'cpu'
+            )
+            vcpu_limit = cls._parse_resource_quantity(
+                limits.get('cpu', limits.get('vcpu')), 'cpu'
+            )
+            memory = cls._parse_resource_quantity(requests.get('memory'), 'memory')
+            memory_limit = cls._parse_resource_quantity(limits.get('memory'), 'memory')
 
             accelerators_limit = None
             if limits.get('accelerators'):

src/sagemaker/hyperpod/training/quota_allocation_util.py

@@ -10,6 +10,7 @@
 # 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.
+from decimal import Decimal
 import re
 from sagemaker.hyperpod.cli.constants.command_constants import NVIDIA_GPU_RESOURCE_LIMIT_KEY, NEURON_RESOURCE_LIMIT_KEY
 from sagemaker.hyperpod.cli.utils import (
@@ -121,7 +122,7 @@ def _trim_resource_requests(instance_type: str, requests_values: dict) -> dict:
     max_allocatable_memory = memory_capacity - (_calculate_memory_reservation(memory_capacity))
 
     cpu_request_str = requests_values.get('cpu', '0')
-    cpu_request = float(''.join(filter(lambda x: x.isdigit() or x == '.', cpu_request_str)))
+    cpu_request = float(cpu_request_str)
 
     mem_request_str = requests_values.get('memory', '0Gi')
     mem_request = float(mem_request_str.replace('Gi', ''))
@@ -185,9 +186,10 @@ def _resolve_default_memory_values(instance_type: str, requests_values: dict, li
     if mem_limit_str is None and mem_request_str is not None:
         mem_limit_str = mem_request_str
 
+    numeric_pattern = r'^([0-9]*\.?[0-9]+(?:[eE][+-]?[0-9]+)?)'
     try:
-        memory_limit = float(re.match(r'^([0-9]*\.?[0-9]+)', mem_limit_str).group(1))
-        memory_request = float(re.match(r'^([0-9]*\.?[0-9]+)', mem_request_str).group(1))
+        memory_limit = float(re.match(numeric_pattern, mem_limit_str).group(1))
+        memory_request = float(re.match(numeric_pattern, mem_request_str).group(1))
     except (AttributeError, ValueError):
         raise ValueError(f"Invalid memory format: {mem_limit_str or mem_request_str}")
 
@@ -203,8 +205,8 @@ def _resolve_default_memory_values(instance_type: str, requests_values: dict, li
         memory_limit = min(memory_limit, max_allocatable_memory)
 
     memory_request = min(memory_request, max_allocatable_memory)
-    limits_values["memory"] = str(memory_limit) + "Gi"
-    requests_values["memory"] = str(memory_request) + "Gi"
+    limits_values["memory"] = format(Decimal(str(memory_limit)), 'f') + "Gi"
+    requests_values["memory"] = format(Decimal(str(memory_request)), 'f') + "Gi"
 
 
 def _validate_accelerators_inputs(instance_type: str, accelerators_request: int, accelerators_limit: int) -> None:

test/unit_tests/training/test_training_resource_quantities.py

@@ -0,0 +1,129 @@
+"""Resource quantities retain their meaning through public job allocation."""
+
+from copy import deepcopy
+
+from hyperpod_pytorch_job_template.v1_1.model import PyTorchJobConfig
+from kubernetes.utils.quantity import parse_quantity
+import pytest
+
+from sagemaker.hyperpod.training import HyperPodPytorchJob
+
+
+def make_job(requests, limits):
+    job = PyTorchJobConfig(
+        job_name="resource-quantities",
+        image="pytorch:latest",
+        instance_type="ml.g5.xlarge",
+        vcpu=1.0,
+        memory=1.0,
+    ).to_domain()
+    resources = job.replicaSpecs[0].template.spec.containers[0].resources
+    resources.requests = requests
+    resources.limits = limits
+    return job
+
+
+@pytest.mark.parametrize(
+    "cpu,memory,cpu_limit,memory_limit,expected_memory,expected_limit",
+    [
+        ("0.5", "1Gi", "1", "2Gi", 2**30, 2**31),
+        ("500m", "1Gi", "1000m", "2Gi", 2**30, 2**31),
+        ("0.5", "1024Mi", "1", "2048Mi", 2**30, 2**31),
+        ("500m", "1024Mi", "1000m", "2048Mi", 2**30, 2**31),
+        ("0.5", "1G", "1", "2G", 10**9, 2 * 10**9),
+        ("0.5", "1e-5", "1", "2e-5", 1e-5 * 2**30, 2e-5 * 2**30),
+        ("0.5", 1, "1", 2, 2**30, 2**31),
+        ("0.5", "1", "1", "2", 2**30, 2**31),
+        ("0.5", "1Ki", "1", "2Ki", 1024, 2048),
+        ("0.5", "1m", "1", "2m", 0.001, 0.002),
+    ],
+)
+def test_standard_quantities_preserved(
+    cpu, memory, cpu_limit, memory_limit, expected_memory, expected_limit
+):
+    job = make_job(
+        {"cpu": cpu, "memory": memory},
+        {"cpu": cpu_limit, "memory": memory_limit},
+    )
+    resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job)
+    resources = resolved.replicaSpecs[0].template.spec.containers[0].resources
+    assert parse_quantity(resources.requests["cpu"]) == parse_quantity("0.5")
+    assert parse_quantity(resources.limits["cpu"]) == 1
+    assert float(parse_quantity(resources.requests["memory"])) == pytest.approx(
+        expected_memory
+    )
+    assert float(parse_quantity(resources.limits["memory"])) == pytest.approx(
+        expected_limit
+    )
+
+
+@pytest.mark.parametrize("section", ["requests", "limits"])
+@pytest.mark.parametrize("resource", ["cpu", "memory"])
+@pytest.mark.parametrize(
+    "invalid", ["NaN", "Infinity", "-1", "1GiB", "1e10000", "1e999999999G"]
+)
+def test_invalid_quantity_fails_before_mutating_job(section, resource, invalid):
+    job = make_job(
+        {"cpu": "0.5", "memory": "1Gi"},
+        {"cpu": "1", "memory": "2Gi"},
+    )
+    resources = job.replicaSpecs[0].template.spec.containers[0].resources
+    getattr(resources, section)[resource] = invalid
+    original = deepcopy(job.model_dump())
+    with pytest.raises(ValueError, match=f"Invalid {resource} quantity"):
+        HyperPodPytorchJob.allocate_quotas_if_applicable(job)
+    assert job.model_dump() == original
+
+
+def test_vcpu_alias_accepts_millicores():
+    job = make_job(
+        {"vcpu": "500m", "memory": "1Gi"},
+        {"vcpu": "1000m", "memory": "2Gi"},
+    )
+    resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job)
+    resources = resolved.replicaSpecs[0].template.spec.containers[0].resources
+    assert parse_quantity(resources.requests["cpu"]) == parse_quantity("0.5")
+    assert parse_quantity(resources.limits["cpu"]) == 1
+
+
+@pytest.mark.parametrize("memory,memory_limit", [(64.0, 128.0), (1e-5, 2e-5)])
+def test_installed_template_memory_stays_in_gib(memory, memory_limit):
+    job = PyTorchJobConfig(
+        job_name="template-memory",
+        image="pytorch:latest",
+        instance_type="ml.p4d.24xlarge",
+        vcpu=16.0,
+        memory=memory,
+        memory_limit=memory_limit,
+    ).to_domain()
+    resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job)
+    resources = resolved.replicaSpecs[0].template.spec.containers[0].resources
+    assert float(parse_quantity(resources.requests["memory"])) == pytest.approx(
+        memory * 2**30
+    )
+    assert float(parse_quantity(resources.limits["memory"])) == pytest.approx(
+        memory_limit * 2**30
+    )
+    assert "e" not in resources.requests["memory"].lower()
+    assert "e" not in resources.limits["memory"].lower()
+
+
+def test_memory_only_quantity_preserves_derived_cpu_exponent():
+    job = make_job({"memory": "1Ki"}, {})
+    resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job)
+    resources = resolved.replicaSpecs[0].template.spec.containers[0].resources
+    expected_cpu = 1024 / (16 * 2**30) * 4
+    assert float(parse_quantity(resources.requests["cpu"])) == pytest.approx(
+        expected_cpu
+    )
+    assert float(parse_quantity(resources.requests["memory"])) == pytest.approx(1024)
+    assert float(parse_quantity(resources.limits["memory"])) == pytest.approx(1024)
+
+
+def test_numeric_zero_preserves_existing_default_allocation():
+    job = make_job({"cpu": 0, "memory": 0}, {"memory": 0})
+    resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job)
+    resources = resolved.replicaSpecs[0].template.spec.containers[0].resources
+    assert float(parse_quantity(resources.requests["cpu"])) == 3
+    assert float(parse_quantity(resources.requests["memory"])) == 11 * 2**30
+    assert resources.limits["memory"] == resources.requests["memory"]