chore: add js support for copies resource requirement, enforce coupling with ResourceRequirements class

Author: evakravi

src/sagemaker/jumpstart/artifacts/resource_requirements.py

@@ -28,6 +28,18 @@
 from sagemaker.session import Session
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 
+REQUIREMENT_TYPE_TO_SPEC_FIELD_NAME_TO_RESOURCE_REQUIREMENT_NAME_MAP: Dict[str, Dict[str, str]] = {
+    "requests": {
+        "num_accelerators": ("num_accelerators", "num_accelerators"),
+        "num_cpus": ("num_cpus", "num_cpus"),
+        "copies": ("copies", "copy_count"),
+        "min_memory_mb": ("memory", "min_memory"),
+    },
+    "limits": {
+        "max_memory_mb": ("memory", "max_memory"),
+    },
+}
+
 
 def _retrieve_default_resources(
     model_id: str,
@@ -113,16 +125,22 @@ def _retrieve_default_resources(
     }
 
     if is_dynamic_container_deployment_supported:
-        requests = {}
-        if "num_accelerators" in default_resource_requirements:
-            requests["num_accelerators"] = default_resource_requirements["num_accelerators"]
-        if "min_memory_mb" in default_resource_requirements:
-            requests["memory"] = default_resource_requirements["min_memory_mb"]
-        if "num_cpus" in default_resource_requirements:
-            requests["num_cpus"] = default_resource_requirements["num_cpus"]
-
-        limits = {}
-        if "max_memory_mb" in default_resource_requirements:
-            limits["memory"] = default_resource_requirements["max_memory_mb"]
-        return ResourceRequirements(requests=requests, limits=limits)
+
+        all_resource_requirement_kwargs = {}
+
+        for (
+            requirement_type,
+            spec_field_to_resource_requirement_map,
+        ) in REQUIREMENT_TYPE_TO_SPEC_FIELD_NAME_TO_RESOURCE_REQUIREMENT_NAME_MAP.items():
+            requirement_type
+            requirement_kwargs = {}
+            for spec_field, resource_requirement in spec_field_to_resource_requirement_map.items():
+                if spec_field in default_resource_requirements:
+                    requirement_kwargs[resource_requirement[0]] = default_resource_requirements[
+                        spec_field
+                    ]
+
+            all_resource_requirement_kwargs[requirement_type] = requirement_kwargs
+
+        return ResourceRequirements(**all_resource_requirement_kwargs)
     return None

src/sagemaker/resource_requirements.py

@@ -16,6 +16,7 @@
 
 import logging
 from typing import Optional
+from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 
 from sagemaker.jumpstart import utils as jumpstart_utils
 from sagemaker.jumpstart import artifacts
@@ -34,7 +35,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     instance_type: Optional[str] = None,
-) -> str:
+) -> ResourceRequirements:
     """Retrieves the default resource requirements for the model matching the given arguments.
 
     Args:

tests/unit/sagemaker/resource_requirements/jumpstart/test_resource_requirements.py

@@ -18,6 +18,10 @@
 import pytest
 
 from sagemaker import resource_requirements
+from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
+from sagemaker.jumpstart.artifacts.resource_requirements import (
+    REQUIREMENT_TYPE_TO_SPEC_FIELD_NAME_TO_RESOURCE_REQUIREMENT_NAME_MAP,
+)
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
 
@@ -129,3 +133,18 @@ def test_jumpstart_no_supported_resource_requirements(patched_get_model_specs):
         resource_requirements.retrieve_default(
             region=region, model_id=model_id, model_version=model_version, scope="training"
         )
+
+
+def test_jumpstart_supports_all_resource_requirement_fields():
+
+    all_tracked_resource_requirement_fields = {
+        field
+        for requirements in REQUIREMENT_TYPE_TO_SPEC_FIELD_NAME_TO_RESOURCE_REQUIREMENT_NAME_MAP.values()
+        for _, field in requirements.values()
+    }
+
+    excluded_resource_requirement_fields = {"requests", "limits"}
+    assert (
+        set(ResourceRequirements().__dict__.keys()) - excluded_resource_requirement_fields
+        == all_tracked_resource_requirement_fields
+    )