From 7d6b205067822b486b417e8c8a9d92d7b734cb6b Mon Sep 17 00:00:00 2001 From: "Jaehyup (Jay) Shin" Date: Sun, 25 Oct 2026 09:06:24 -0700 Subject: [PATCH] fix: handle resource units during training quota allocation Preserve legacy unitless GiB inputs while normalizing explicit Kubernetes units. Keep exponent-form quantities intact through allocation and reject invalid values before mutation. Signed-off-by: shin0343 --- doc/getting_started/training.md | 14 ++ .../hyperpod/training/hyperpod_pytorch_job.py | 62 ++++++--- .../training/quota_allocation_util.py | 12 +- .../test_training_resource_quantities.py | 129 ++++++++++++++++++ 4 files changed, 191 insertions(+), 26 deletions(-) create mode 100644 test/unit_tests/training/test_training_resource_quantities.py diff --git a/doc/getting_started/training.md b/doc/getting_started/training.md index 4bbdf42e..e1793650 100644 --- a/doc/getting_started/training.md +++ b/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: diff --git a/src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py b/src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py index f39bc07e..a43b68fc 100644 --- a/src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py +++ b/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'): diff --git a/src/sagemaker/hyperpod/training/quota_allocation_util.py b/src/sagemaker/hyperpod/training/quota_allocation_util.py index 1b9df4a6..e8f3fd99 100644 --- a/src/sagemaker/hyperpod/training/quota_allocation_util.py +++ b/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: diff --git a/test/unit_tests/training/test_training_resource_quantities.py b/test/unit_tests/training/test_training_resource_quantities.py new file mode 100644 index 00000000..c9302657 --- /dev/null +++ b/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"]