feat: JumpStart alternative config parsing

src/sagemaker/jumpstart/enums.py

@@ -124,3 +124,9 @@ def from_suffixed_type(mime_type_with_suffix: str) -> "MIMEType":
         """Removes suffix from type and instantiates enum."""
         base_type, _, _ = mime_type_with_suffix.partition(";")
         return MIMEType(base_type)
+
+
+class JumpStartConfigRankingName(str, Enum):
+    """Enum class for ranking of JumpStart config."""
+
+    DEFAULT = "overall"

src/sagemaker/jumpstart/utils.py

@@ -37,9 +37,9 @@
 )
 from sagemaker.jumpstart.types import (
     JumpStartBenchmarkStat,
+    JumpStartMetadataConfig,
     JumpStartModelHeader,
     JumpStartModelSpecs,
-    JumpStartPresetConfig,
     JumpStartVersionedModelId,
 )
 from sagemaker.session import Session
@@ -882,15 +882,15 @@ def get_region_fallback(
     return list(combined_regions)[0]
 
 
-def get_preset_names(
+def get_config_names(
     region: str,
     model_id: str,
     model_version: str,
     sagemaker_session: Optional[Session] = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     scope: enums.JumpStartScriptScope = enums.JumpStartScriptScope.INFERENCE,
     model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS,
 ) -> List[str]:
-    """Returns a list of preset names for the given model ID and region."""
+    """Returns a list of config names for the given model ID and region."""
     model_specs = verify_model_region_and_return_specs(
         region=region,
         model_id=model_id,
@@ -901,12 +901,13 @@ def get_preset_names(
     )
 
     if scope == enums.JumpStartScriptScope.INFERENCE:
-        presets = model_specs.inference_presets
-
-    if scope == enums.JumpStartScriptScope.TRAINING:
-        presets = model_specs.training_presets
+        metadata_configs = model_specs.inference_configs
+    elif scope == enums.JumpStartScriptScope.TRAINING:
+        metadata_configs = model_specs.training_configs
+    else:
+        raise ValueError(f"Unknown script scope {scope}.")
 
-    return list(presets.preset_configs.keys()) if presets else []
+    return list(metadata_configs.configs.keys()) if metadata_configs else []
 
 
 def get_benchmark_stats(
@@ -929,31 +930,34 @@ def get_benchmark_stats(
     )
 
     if scope == enums.JumpStartScriptScope.INFERENCE:
-        presets = model_specs.inference_presets
+        metadata_configs = model_specs.inference_configs
     elif scope == enums.JumpStartScriptScope.TRAINING:
-        presets = model_specs.training_presets
+        metadata_configs = model_specs.training_configs
+    else:
+        raise ValueError(f"Unknown script scope {scope}.")
 
     if not config_names:
-        config_names = presets.preset_configs.keys() if presets else []
+        config_names = metadata_configs.configs.keys() if metadata_configs else []
 
     benchmark_stats = {}
     for config_name in config_names:
-        if config_name not in presets.preset_configs:
-            raise ValueError(f"Unknown preset config name: '{config_name}'")
-        benchmark_stats[config_name] = presets.preset_configs.get(config_name).benchmark_metrics
+        if config_name not in metadata_configs.configs:
+            raise ValueError(f"Unknown config name: '{config_name}'")
+        benchmark_stats[config_name] = metadata_configs.configs.get(config_name).benchmark_metrics
 
     return benchmark_stats
 
 
-def get_jumpstart_presets(
+def get_jumpstart_configs(
     region: str,
     model_id: str,
     model_version: str,
     config_names: Optional[List[str]] = None,
     sagemaker_session: Optional[Session] = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     scope: enums.JumpStartScriptScope = enums.JumpStartScriptScope.INFERENCE,
     model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS,
-) -> Dict[str, List[JumpStartPresetConfig]]:
+) -> Dict[str, List[JumpStartMetadataConfig]]:
+    """Returns metadata configs for the given model ID and region."""
     model_specs = verify_model_region_and_return_specs(
         region=region,
         model_id=model_id,
@@ -964,15 +968,17 @@ def get_jumpstart_presets(
     )
 
     if scope == enums.JumpStartScriptScope.INFERENCE:
-        presets = model_specs.inference_presets
+        metadata_configs = model_specs.inference_configs
     elif scope == enums.JumpStartScriptScope.TRAINING:
-        presets = model_specs.training_presets
+        metadata_configs = model_specs.training_configs
+    else:
+        raise ValueError(f"Unknown script scope {scope}.")
 
     if not config_names:
-        config_names = presets.preset_configs.keys()
+        config_names = metadata_configs.configs.keys() if metadata_configs else []
 
-    preset_configs = {
-        config_name: presets.preset_configs[config_name] for config_name in config_names
-    }
-
-    return preset_configs
+    return (
+        {config_name: metadata_configs.configs[config_name] for config_name in config_names}
+        if metadata_configs
+        else {}
+    )

src/sagemaker/utils.py

@@ -1602,36 +1602,36 @@ def can_model_package_source_uri_autopopulate(source_uri: str):
     )
 
 
-def flatten_dict(dict: Dict[str, Any], sep: str = ".") -> Dict[str, Any]:
+def flatten_dict(source_dict: Dict[str, Any], sep: str = ".") -> Dict[str, Any]:
     """Flatten a nested dictionary.
+
     Args:
-        dict (dict): The dictionary to be flattened.
+        source_dict (dict): The dictionary to be flattened.
         sep (str): The separator to be used in the flattened dictionary.
     Returns:
-        dict: The flattened dictionary.
+        transformed_dict: The flattened dictionary.
     """
-    flat_dict_list = pd.json_normalize(dict, sep=sep).to_dict(orient="records")
+    flat_dict_list = pd.json_normalize(source_dict, sep=sep).to_dict(orient="records")
     if flat_dict_list:
         return flat_dict_list[0]
     return {}
 
 
-def unflatten_dict(dict: Dict[str, Any], sep: str = ".") -> Dict[str, Any]:
-    """
-    Unflatten a flattened dictionary back into a nested dictionary.
+def unflatten_dict(source_dict: Dict[str, Any], sep: str = ".") -> Dict[str, Any]:
+    """Unflatten a flattened dictionary back into a nested dictionary.
 
     Args:
-        d (dict): The input flattened dictionary.
+        source_dict (dict): The input flattened dictionary.
         sep (str): The separator used in the flattened keys.
 
     Returns:
-        dict: The reconstructed nested dictionary.
+        transformed_dict: The reconstructed nested dictionary.
     """
-    if not dict:
+    if not source_dict:
         return {}
 
     result = {}
-    for key, value in dict.items():
+    for key, value in source_dict.items():
         keys = key.split(sep)
         current = result
         for k in keys[:-1]:
@@ -1646,6 +1646,9 @@ def deep_override_dict(
     dict1: Dict[str, Any], dict2: Dict[str, Any], skip_keys: Optional[List[str]] = None
 ) -> Dict[str, Any]:
     """Overrides any overlapping contents of dict1 with the contents of dict2."""
+    if skip_keys is None:
+        skip_keys = []
+
     flattened_dict1 = flatten_dict(dict1)
     flattened_dict2 = flatten_dict(
         {key: value for key, value in dict2.items() if key not in skip_keys}

tests/unit/sagemaker/jumpstart/constants.py

@@ -7505,12 +7505,12 @@
     "resource_name_base": "dfsdfsds",
     "hosting_resource_requirements": {"num_accelerators": 1, "min_memory_mb": 34360},
     "dynamic_container_deployment_supported": True,
-    "inference_presets": None,
-    "inference_preset_components": None,
-    "training_presets": None,
-    "training_preset_components": None,
-    "inference_preset_rankings": None,
-    "training_preset_rankings": None,
+    "inference_configs": None,
+    "inference_config_components": None,
+    "training_configs": None,
+    "training_config_components": None,
+    "inference_config_rankings": None,
+    "training_config_rankings": None,
 }
 
 BASE_HEADER = {
@@ -7654,17 +7654,17 @@
 }
 
 
-INFERENCE_PRESETS = {
-    "inference_presets": {
+INFERENCE_CONFIGS = {
+    "inference_configs": {
         "neuron-inference": {
             "benchmark_metrics": {
-                "ml.p3.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.inf2.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
             },
             "component_names": ["neuron-base"],
         },
         "neuron-inference-budget": {
             "benchmark_metrics": {
-                "ml.p3.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.inf2.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
             },
             "component_names": ["neuron-base"],
         },
@@ -7681,7 +7681,7 @@
             "component_names": ["gpu-inference"],
         },
     },
-    "inference_preset_components": {
+    "inference_config_components": {
         "neuron-base": {
             "supported_inference_instance_types": ["ml.inf2.xlarge", "ml.inf2.2xlarge"]
         },
@@ -7734,8 +7734,8 @@
     },
 }
 
-TRAINING_PRESETS = {
-    "training_presets": {
+TRAINING_CONFIGS = {
+    "training_configs": {
         "neuron-training": {
             "benchmark_metrics": {
                 "ml.tr1n1.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"},
@@ -7763,7 +7763,7 @@
             "component_names": ["gpu-training-budget"],
         },
     },
-    "training_preset_components": {
+    "training_config_components": {
         "neuron-training": {
             "supported_training_instance_types": ["ml.trn1.xlarge", "ml.trn1.2xlarge"],
             "training_artifact_key": "artifacts/meta-textgeneration-llama-2-7b/neuron-training/model/",
@@ -7826,70 +7826,74 @@
 }
 
 
-INFERENCE_PRESET_RANKINGS = {
-    "overall": {
-        "description": "Overall rankings of configs",
-        "ranking": [
-            "neuron-inference",
-            "neuron-inference-budget",
-            "gpu-inference",
-            "gpu-inference-budget",
-        ],
-    },
-    "performance": {
-        "description": "Configs ranked based on performance",
-        "rankings": [
-            "neuron-inference",
-            "gpu-inference",
-            "neuron-inference-budget",
-            "gpu-inference-budget",
-        ],
-    },
-    "cost": {
-        "description": "Configs ranked based on cost",
-        "rankings": [
-            "neuron-inference-budget",
-            "gpu-inference-budget",
-            "neuron-inference",
-            "gpu-inference",
-        ],
-    },
+INFERENCE_CONFIG_RANKINGS = {
+    "inference_config_rankings": {
+        "overall": {
+            "description": "Overall rankings of configs",
+            "rankings": [
+                "neuron-inference",
+                "neuron-inference-budget",
+                "gpu-inference",
+                "gpu-inference-budget",
+            ],
+        },
+        "performance": {
+            "description": "Configs ranked based on performance",
+            "rankings": [
+                "neuron-inference",
+                "gpu-inference",
+                "neuron-inference-budget",
+                "gpu-inference-budget",
+            ],
+        },
+        "cost": {
+            "description": "Configs ranked based on cost",
+            "rankings": [
+                "neuron-inference-budget",
+                "gpu-inference-budget",
+                "neuron-inference",
+                "gpu-inference",
+            ],
+        },
+    }
 }
 
-TRAINING_PRESET_RANKINGS = {
-    "overall": {
-        "description": "Overall rankings of configs",
-        "rankings": [
-            "neuron-training",
-            "neuron-training-budget",
-            "gpu-training",
-            "gpu-training-budget",
-        ],
-    },
-    "performance_training": {
-        "description": "Configs ranked based on performance",
-        "rankings": [
-            "neuron-training",
-            "gpu-training",
-            "neuron-training-budget",
-            "gpu-training-budget",
-        ],
-        "instance_type_overrides": {
-            "ml.p2.xlarge": [
+TRAINING_CONFIG_RANKINGS = {
+    "training_config_rankings": {
+        "overall": {
+            "description": "Overall rankings of configs",
+            "rankings": [
                 "neuron-training",
                 "neuron-training-budget",
                 "gpu-training",
                 "gpu-training-budget",
-            ]
+            ],
         },
-    },
-    "cost_training": {
-        "description": "Configs ranked based on cost",
-        "rankings": [
-            "neuron-training-budget",
-            "gpu-training-budget",
-            "neuron-training",
-            "gpu-training",
-        ],
-    },
+        "performance_training": {
+            "description": "Configs ranked based on performance",
+            "rankings": [
+                "neuron-training",
+                "gpu-training",
+                "neuron-training-budget",
+                "gpu-training-budget",
+            ],
+            "instance_type_overrides": {
+                "ml.p2.xlarge": [
+                    "neuron-training",
+                    "neuron-training-budget",
+                    "gpu-training",
+                    "gpu-training-budget",
+                ]
+            },
+        },
+        "cost_training": {
+            "description": "Configs ranked based on cost",
+            "rankings": [
+                "neuron-training-budget",
+                "gpu-training-budget",
+                "neuron-training",
+                "gpu-training",
+            ],
+        },
+    }
 }