diff --git a/doc/getting_started/training.md b/doc/getting_started/training.md index 4bbdf42e..c9d7b4e6 100644 --- a/doc/getting_started/training.md +++ b/doc/getting_started/training.md @@ -267,6 +267,16 @@ pytorch_job.create() ```` ````` +For SDK jobs with several `ReplicaSpec` entries, quota allocation validates and +resolves the main container's resources in every replica set using that replica's +instance type. For example, a chief on `ml.g5.xlarge` and workers on +`ml.p4d.24xlarge` receive their own GPU, CPU, memory, and EFA allocations. Replica +sets without an instance type keep their explicit resources, including omitted +resource mappings. Managed replicas can omit mappings to use instance defaults. +If any replica set +fails validation, allocation raises `ValueError` without partially updating the +input job. + ### 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..dd5e74ce 100644 --- a/src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py +++ b/src/sagemaker/hyperpod/training/hyperpod_pytorch_job.py @@ -7,7 +7,7 @@ EFA_RESOURCE_LIMIT_KEY, ) from sagemaker.hyperpod.training.config.hyperpod_pytorch_job_unified_config import ( - _HyperPodPytorchJob, HyperPodPytorchJobStatus + _HyperPodPytorchJob, HyperPodPytorchJobStatus, Resources ) from sagemaker.hyperpod.common.config.metadata import Metadata from kubernetes import client, config, stream @@ -100,20 +100,20 @@ def _process_replica_resources(cls, data): node_count = data.get('replicas', None) # Extract nested configuration with validation - template = data.get('template', {}) - spec = template.get('spec', {}) + template = data.get('template') or {} + spec = template.get('spec') or {} node_selector = spec.get('nodeSelector', {}) instance_type = node_selector.get(INSTANCE_TYPE_LABEL) if node_selector else None - containers = spec.get('containers', []) + containers = spec.get('containers') or [] 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', {}) + resources = container.get('resources') or {} + requests = resources.get('requests') or {} + limits = resources.get('limits') or {} # Check if accelerators are specified without instance_type accel_keys = ['accelerators', NVIDIA_RESOURCE_KEY, NEURON_RESOURCE_KEY] @@ -200,8 +200,9 @@ def _process_replica_resources(cls, data): _resolve_default_cpu_values(instance_type, requests_values) # Update data with calculated values - data['template']['spec']['containers'][0]['resources']['requests'] = requests_values - data['template']['spec']['containers'][0]['resources']['limits'] = limits_values + resources['requests'] = requests_values + resources['limits'] = limits_values + container['resources'] = resources return data except KeyError as e: @@ -215,22 +216,29 @@ def _get_container_resources(cls, replica_spec): @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) + """Resolve every replica set's resources before updating the job. - # Update the original spec object directly + Replica sets without an instance type keep their explicit resources. + A validation failure in any replica set leaves the input job unchanged. + """ + spec_dict = spec.model_dump() + allocations = [] + for index, replica_spec in enumerate(spec_dict['replicaSpecs'] or []): + if cls._process_replica_resources(replica_spec) is None: + continue 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 + allocations.append((index, requests, limits)) - return spec - except ValueError as e: - raise ValueError(e) - except Exception as e: - # In case of any other exception, return original spec - return spec + # Validate all replica sets before mutating the caller's job. + for index, requests, limits in allocations: + container = spec.replicaSpecs[index].template.spec.containers[0] + if container.resources is None: + container.resources = Resources(requests=requests, limits=limits) + else: + container.resources.requests = requests + container.resources.limits = limits + + return spec @_hyperpod_telemetry_emitter(Feature.HYPERPOD, "create_pytorchjob") def create(self, debug=False): diff --git a/test/unit_tests/training/test_multi_replica_resource_allocation.py b/test/unit_tests/training/test_multi_replica_resource_allocation.py new file mode 100644 index 00000000..042bacc2 --- /dev/null +++ b/test/unit_tests/training/test_multi_replica_resource_allocation.py @@ -0,0 +1,182 @@ +"""Every training replica set receives validation and quota allocation.""" + +from copy import deepcopy +from unittest.mock import patch + +from hyperpod_pytorch_job_template.v1_1.model import PyTorchJobConfig +import pytest + +from sagemaker.hyperpod.training import HyperPodPytorchJob, Resources + + +def make_multi_replica_job(): + chief = PyTorchJobConfig( + job_name="multi-replica", + image="pytorch:latest", + instance_type="ml.g5.xlarge", + accelerators=1, + ).to_domain() + worker = PyTorchJobConfig( + job_name="worker", + image="pytorch:latest", + instance_type="ml.p4d.24xlarge", + accelerators=4, + replica_count=2, + ).to_domain() + chief.replicaSpecs[0].name = "chief" + worker.replicaSpecs[0].name = "workers" + chief.replicaSpecs.extend(worker.replicaSpecs) + return chief + + +def test_all_replica_sets_receive_instance_specific_resources(): + job = make_multi_replica_job() + resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job) + chief, workers = resolved.replicaSpecs + chief_resources = chief.template.spec.containers[0].resources + worker_resources = workers.template.spec.containers[0].resources + + assert int(chief_resources.requests["nvidia.com/gpu"]) == 1 + assert "vpc.amazonaws.com/efa" not in chief_resources.requests + assert int(worker_resources.requests["nvidia.com/gpu"]) == 4 + assert int(worker_resources.limits["nvidia.com/gpu"]) == 4 + assert int(worker_resources.requests["vpc.amazonaws.com/efa"]) == 4 + assert int(worker_resources.limits["vpc.amazonaws.com/efa"]) == 4 + assert "accelerators" not in worker_resources.requests + assert [replica.name for replica in resolved.replicaSpecs] == ["chief", "workers"] + assert workers.replicas == 2 + + +@pytest.mark.parametrize("unmanaged_index", [0, 1]) +def test_replica_without_instance_type_does_not_skip_other_replicas(unmanaged_index): + job = make_multi_replica_job() + unmanaged = job.replicaSpecs[unmanaged_index] + unmanaged.template.spec.nodeSelector = None + unmanaged.template.spec.containers[0].resources.requests = {"cpu": "1"} + unmanaged.template.spec.containers[0].resources.limits = {"cpu": "2"} + original = deepcopy(unmanaged.model_dump()) + + resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job) + assert resolved.replicaSpecs[unmanaged_index].model_dump() == original + managed = resolved.replicaSpecs[1 - unmanaged_index] + resources = managed.template.spec.containers[0].resources + assert "nvidia.com/gpu" in resources.requests + assert "accelerators" not in resources.requests + + +@pytest.mark.parametrize("invalid", ["accelerator-capacity", "missing-instance-type"]) +def test_later_invalid_replica_is_rejected_without_partial_mutation(invalid): + job = make_multi_replica_job() + worker = job.replicaSpecs[1] + if invalid == "accelerator-capacity": + resources = worker.template.spec.containers[0].resources + resources.requests["accelerators"] = "9" + resources.limits["accelerators"] = "9" + expected_error = "exceeds instance capacity" + else: + worker.template.spec.nodeSelector = None + expected_error = "instance_type is required" + original = deepcopy(job.model_dump()) + + with pytest.raises(ValueError, match=expected_error): + HyperPodPytorchJob.allocate_quotas_if_applicable(job) + assert job.model_dump() == original + + +@pytest.mark.parametrize("index", [0, 1]) +@pytest.mark.parametrize( + "resources", + [ + None, + Resources(), + Resources(requests={"cpu": "1"}), + Resources(limits={"memory": "8Gi"}), + ], +) +def test_unmanaged_replica_with_optional_resources_preserves_other_allocations( + index, resources +): + job = make_multi_replica_job() + unmanaged = job.replicaSpecs[index] + unmanaged.template.spec.nodeSelector = None + unmanaged.template.spec.containers[0].resources = resources + original = deepcopy(unmanaged.model_dump()) + resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job) + assert resolved.replicaSpecs[index].model_dump() == original + managed_resources = ( + resolved.replicaSpecs[1 - index].template.spec.containers[0].resources + ) + assert "nvidia.com/gpu" in managed_resources.requests + assert "accelerators" not in managed_resources.requests + + +@pytest.mark.parametrize( + "resources", + [ + None, + Resources(), + Resources(requests={"cpu": "1"}), + Resources(limits={"memory": "8Gi"}), + ], +) +def test_managed_replica_with_optional_resources_is_allocated(resources): + job = make_multi_replica_job() + job.replicaSpecs[1].template.spec.containers[0].resources = resources + resolved = HyperPodPytorchJob.allocate_quotas_if_applicable(job) + chief_resources = resolved.replicaSpecs[0].template.spec.containers[0].resources + worker_resources = resolved.replicaSpecs[1].template.spec.containers[0].resources + assert int(chief_resources.requests["nvidia.com/gpu"]) == 1 + assert "cpu" in worker_resources.requests + assert "memory" in worker_resources.requests + assert "accelerators" not in worker_resources.requests + + +@patch("sagemaker.hyperpod.common.telemetry.telemetry_logging._send_telemetry_request") +@patch("sagemaker.hyperpod.training.hyperpod_pytorch_job.client.CustomObjectsApi") +@patch.object(HyperPodPytorchJob, "verify_kube_config") +def test_create_serializes_allocations_with_optional_worker_resources( + verify_config, api, telemetry +): + job = make_multi_replica_job() + job.metadata.namespace = "default" + worker = job.replicaSpecs[1] + worker.template.spec.nodeSelector = None + worker.template.spec.containers[0].resources = None + job.create() + body = api.return_value.create_namespaced_custom_object.call_args.kwargs["body"] + chief, worker_payload = body["spec"]["replicaSpecs"] + requests = chief["template"]["spec"]["containers"][0]["resources"]["requests"] + assert int(requests["nvidia.com/gpu"]) == 1 + assert "cpu" in requests and "memory" in requests + assert "resources" not in worker_payload["template"]["spec"]["containers"][0] + + +@patch("sagemaker.hyperpod.common.telemetry.telemetry_logging._send_telemetry_request") +@patch("sagemaker.hyperpod.training.hyperpod_pytorch_job.client.CustomObjectsApi") +@patch.object(HyperPodPytorchJob, "verify_kube_config") +def test_create_rejects_invalid_later_replica_before_api_submission( + verify_config, api, telemetry +): + job = make_multi_replica_job() + job.metadata.namespace = "default" + job.replicaSpecs[1].template.spec.nodeSelector = None + with pytest.raises(ValueError, match="instance_type is required"): + job.create() + api.assert_not_called() + + +@pytest.mark.parametrize("replicas", [None, []]) +def test_missing_replica_list_remains_unchanged(replicas): + job = make_multi_replica_job() + job.replicaSpecs = replicas + assert HyperPodPytorchJob.allocate_quotas_if_applicable(job) is job + assert job.replicaSpecs == replicas + + +def test_later_replica_without_template_fails_without_partial_mutation(): + job = make_multi_replica_job() + job.replicaSpecs[1].template = None + original = deepcopy(job.model_dump()) + with pytest.raises(ValueError, match="No containers found"): + HyperPodPytorchJob.allocate_quotas_if_applicable(job) + assert job.model_dump() == original