feat: Support Alt Configs for Public & Curated Hub (aws#1505)

akozd · web-flow · commit f1bc99e92c65 · 2024-07-03T13:53:22.000-04:00
* feat: add alt config support for public &amp; curated hub
diff --git a/src/sagemaker/jumpstart/hub/interfaces.py b/src/sagemaker/jumpstart/hub/interfaces.py
@@ -13,14 +13,20 @@
 """This module stores types related to SageMaker JumpStart HubAPI requests and responses."""
 from __future__ import absolute_import
 
+from enum import Enum
 import re
 import json
 import datetime
 
 from typing import Any, Dict, List, Union, Optional
+from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.types import (
     HubContentType,
     HubArnExtractedInfo,
+    JumpStartConfigComponent,
+    JumpStartConfigRanking,
+    JumpStartMetadataConfig,
+    JumpStartMetadataConfigs,
     JumpStartPredictorSpecs,
     JumpStartHyperparameter,
     JumpStartDataHolderType,
@@ -34,6 +40,13 @@
 )
 
 
+class _ComponentType(str, Enum):
+    """Enum for different component types."""
+
+    INFERENCE = "Inference"
+    TRAINING = "Training"
+
+
 class HubDataHolderType(JumpStartDataHolderType):
     """Base class for many Hub API interfaces."""
 
@@ -456,6 +469,9 @@ class HubModelDocument(HubDataHolderType):
         "hosting_use_script_uri",
         "hosting_eula_uri",
         "hosting_model_package_arn",
+        "inference_configs",
+        "inference_config_components",
+        "inference_config_rankings",
         "training_artifact_s3_data_type",
         "training_artifact_compression_type",
         "training_model_package_artifact_uri",
@@ -467,6 +483,9 @@ class HubModelDocument(HubDataHolderType):
         "training_ecr_uri",
         "training_metrics",
         "training_artifact_uri",
+        "training_configs",
+        "training_config_components",
+        "training_config_rankings",
         "inference_dependencies",
         "training_dependencies",
         "default_inference_instance_type",
@@ -566,6 +585,11 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         )
         self.hosting_eula_uri: Optional[str] = json_obj.get("HostingEulaUri")
         self.hosting_model_package_arn: Optional[str] = json_obj.get("HostingModelPackageArn")
+
+        self.inference_config_rankings = self._get_config_rankings(json_obj)
+        self.inference_config_components = self._get_config_components(json_obj)
+        self.inference_configs = self._get_configs(json_obj)
+
         self.default_inference_instance_type: Optional[str] = json_obj.get(
             "DefaultInferenceInstanceType"
         )
@@ -667,6 +691,15 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
                 "TrainingMetrics", None
             )
             self.training_artifact_uri: Optional[str] = json_obj.get("TrainingArtifactUri")
+
+            self.training_config_rankings = self._get_config_rankings(
+                json_obj, _ComponentType.TRAINING
+            )
+            self.training_config_components = self._get_config_components(
+                json_obj, _ComponentType.TRAINING
+            )
+            self.training_configs = self._get_configs(json_obj, _ComponentType.TRAINING)
+
             self.training_dependencies: Optional[str] = json_obj.get("TrainingDependencies")
             self.default_training_instance_type: Optional[str] = json_obj.get(
                 "DefaultTrainingInstanceType"
@@ -707,6 +740,64 @@ def get_region(self) -> str:
         """Returns hub region."""
         return self._region
 
+    def _get_config_rankings(
+        self, json_obj: Dict[str, Any], component_type=_ComponentType.INFERENCE
+    ) -> Optional[Dict[str, JumpStartConfigRanking]]:
+        """Returns config rankings."""
+        config_rankings = json_obj.get(f"{component_type.value}ConfigRankings")
+        return (
+            {
+                alias: JumpStartConfigRanking(ranking, is_hub_content=True)
+                for alias, ranking in config_rankings.items()
+            }
+            if config_rankings
+            else None
+        )
+
+    def _get_config_components(
+        self, json_obj: Dict[str, Any], component_type=_ComponentType.INFERENCE
+    ) -> Optional[Dict[str, JumpStartConfigComponent]]:
+        """Returns config components."""
+        config_components = json_obj.get(f"{component_type.value}ConfigComponents")
+        return (
+            {
+                alias: JumpStartConfigComponent(alias, config, is_hub_content=True)
+                for alias, config in config_components.items()
+            }
+            if config_components
+            else None
+        )
+
+    def _get_configs(
+        self, json_obj: Dict[str, Any], component_type=_ComponentType.INFERENCE
+    ) -> Optional[JumpStartMetadataConfigs]:
+        """Returns configs."""
+        if not (configs := json_obj.get(f"{component_type.value}Configs")):
+            return None
+
+        configs_dict = {}
+        for alias, config in configs.items():
+            config_components = None
+            if isinstance(config, dict) and (component_names := config.get("ComponentNames")):
+                config_components = {
+                    name: getattr(self, f"{component_type.value.lower()}_config_components").get(
+                        name
+                    )
+                    for name in component_names
+                }
+            configs_dict[alias] = JumpStartMetadataConfig(
+                alias, config, json_obj, config_components, is_hub_content=True
+            )
+
+        if component_type == _ComponentType.INFERENCE:
+            config_rankings = self.inference_config_rankings
+            scope = JumpStartScriptScope.INFERENCE
+        else:
+            config_rankings = self.training_config_rankings
+            scope = JumpStartScriptScope.TRAINING
+
+        return JumpStartMetadataConfigs(configs_dict, config_rankings, scope)
+
 
 class HubNotebookDocument(HubDataHolderType):
     """Data class for notebook type HubContentDocument from session.describe_hub_content()."""
diff --git a/src/sagemaker/jumpstart/hub/parsers.py b/src/sagemaker/jumpstart/hub/parsers.py
@@ -142,6 +142,9 @@ def make_model_specs_from_describe_hub_content_response(
         hub_model_document.incremental_training_supported
     )
     specs["hosting_ecr_uri"] = hub_model_document.hosting_ecr_uri
+    specs["inference_configs"] = hub_model_document.inference_configs
+    specs["inference_config_components"] = hub_model_document.inference_config_components
+    specs["inference_config_rankings"] = hub_model_document.inference_config_rankings
 
     hosting_artifact_bucket, hosting_artifact_key = parse_s3_url(  # pylint: disable=unused-variable
         hub_model_document.hosting_artifact_uri
@@ -233,6 +236,11 @@ def make_model_specs_from_describe_hub_content_response(
             training_script_key,
         ) = parse_s3_url(hub_model_document.training_script_uri)
         specs["training_script_key"] = training_script_key
+
+        specs["training_configs"] = hub_model_document.training_configs
+        specs["training_config_components"] = hub_model_document.training_config_components
+        specs["training_config_rankings"] = hub_model_document.training_config_rankings
+
         specs["training_dependencies"] = hub_model_document.training_dependencies
         specs["default_training_instance_type"] = hub_model_document.default_training_instance_type
         specs["supported_training_instance_types"] = (
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -1169,12 +1169,14 @@ class JumpStartConfigRanking(JumpStartDataHolderType):
 
     __slots__ = ["description", "rankings"]
 
-    def __init__(self, spec: Optional[Dict[str, Any]]):
+    def __init__(self, spec: Optional[Dict[str, Any]], is_hub_content=False):
         """Initializes a JumpStartConfigRanking object.
 
         Args:
             spec (Dict[str, Any]): Dictionary representation of training config ranking.
         """
+        if is_hub_content:
+            spec = {camel_to_snake(key): val for key, val in spec.items()}
         self.from_json(spec)
 
     def from_json(self, json_obj: Dict[str, Any]) -> None:
@@ -1285,7 +1287,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             json_obj.get("incremental_training_supported", False)
         )
         if self._is_hub_content:
-            self.hosting_ecr_uri: Optional[str] = json_obj["hosting_ecr_uri"]
+            self.hosting_ecr_uri: Optional[str] = json_obj.get("hosting_ecr_uri")
             self._non_serializable_slots.append("hosting_ecr_specs")
         else:
             self.hosting_ecr_specs: Optional[JumpStartECRSpecs] = (
@@ -1491,9 +1493,7 @@ class JumpStartConfigComponent(JumpStartMetadataBaseFields):
     __slots__ = slots + JumpStartMetadataBaseFields.__slots__
 
     def __init__(
-        self,
-        component_name: str,
-        component: Optional[Dict[str, Any]],
+        self, component_name: str, component: Optional[Dict[str, Any]], is_hub_content=False
     ):
         """Initializes a JumpStartConfigComponent object from its json representation.
 
@@ -1504,8 +1504,10 @@ def __init__(
         Raises:
             ValueError: If the component field is invalid.
         """
-        super().__init__(component)
+        if is_hub_content:
+            component = walk_and_apply_json(component, camel_to_snake)
         self.component_name = component_name
+        super().__init__(component, is_hub_content)
         self.from_json(component)
 
     def from_json(self, json_obj: Dict[str, Any]) -> None:
@@ -1542,6 +1544,7 @@ def __init__(
         config: Dict[str, Any],
         base_fields: Dict[str, Any],
         config_components: Dict[str, JumpStartConfigComponent],
+        is_hub_content=False,
     ):
         """Initializes a JumpStartMetadataConfig object from its json representation.
 
@@ -1554,6 +1557,9 @@ def __init__(
             config_components (Dict[str, JumpStartConfigComponent]):
                 The list of components that are used to construct the resolved config.
         """
+        if is_hub_content:
+            config = walk_and_apply_json(config, camel_to_snake)
+            base_fields = walk_and_apply_json(base_fields, camel_to_snake)
         self.base_fields = base_fields
         self.config_components: Dict[str, JumpStartConfigComponent] = config_components
         self.benchmark_metrics: Dict[str, List[JumpStartBenchmarkStat]] = (
diff --git a/tests/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
@@ -8703,7 +8703,17 @@
                 "variants": {"inf2": {"regional_properties": {"image_uri": "$neuron-ecr-uri"}}},
             },
         },
-        "neuron-budget": {"inference_environment_variables": {"BUDGET": "1234"}},
+        "neuron-budget": {
+            "inference_environment_variables": [
+                {
+                    "name": "SAGEMAKER_PROGRAM",
+                    "type": "text",
+                    "default": "inference.py",
+                    "scope": "container",
+                    "required_for_model_class": True,
+                }
+            ],
+        },
         "gpu-inference": {
             "supported_inference_instance_types": ["ml.p2.xlarge", "ml.p3.2xlarge"],
             "hosting_artifact_key": "artifacts/meta-textgeneration-llama-2-7b/gpu-inference/model/",
@@ -9816,6 +9826,58 @@
         "DynamicContainerDeploymentSupported": True,
         "TrainingModelPackageArtifactUri": None,
         "Dependencies": [],
+        "InferenceConfigRankings": {
+            "overall": {"Description": "default", "Rankings": ["variant1"]}
+        },
+        "InferenceConfigs": {
+            "variant1": {
+                "ComponentNames": ["variant1"],
+                "BenchmarkMetrics": {
+                    "ml.g5.12xlarge": [
+                        {"Name": "latency", "Unit": "sec", "Value": "0.19", "Concurrency": "1"},
+                    ]
+                },
+            },
+        },
+        "InferenceConfigComponents": {
+            "variant1": {
+                "HostingEcrUri": "123456789012.ecr.us-west-2.amazon.com/repository",
+                "HostingArtifactUri": "s3://jumpstart-private-cache-prod-us-west-2/meta-textgeneration/meta-textgeneration-llama-2-7b/artifacts/variant1/v1.0.0/",  # noqa: E501
+                "HostingScriptUri": "s3://jumpstart-monarch-test-hub-bucket/monarch-curated-hub-1714579993.88695/curated_models/meta-textgeneration-llama-2-7b/4.0.0/source-directory-tarballs/meta/inference/textgeneration/v1.2.3/sourcedir.tar.gz",  # noqa: E501
+                "InferenceDependencies": [],
+                "InferenceEnvironmentVariables": [
+                    {
+                        "Name": "SAGEMAKER_PROGRAM",
+                        "Type": "text",
+                        "Default": "inference.py",
+                        "Scope": "container",
+                        "RequiredForModelClass": True,
+                    }
+                ],
+                "HostingAdditionalDataSources": {
+                    "speculative_decoding": [
+                        {
+                            "ArtifactVersion": 1,
+                            "ChannelName": "speculative_decoding_channel_1",
+                            "S3DataSource": {
+                                "CompressionType": "None",
+                                "S3DataType": "S3Prefix",
+                                "S3Uri": "s3://bucket/path/1",
+                            },
+                        },
+                        {
+                            "ArtifactVersion": 1,
+                            "ChannelName": "speculative_decoding_channel_2",
+                            "S3DataSource": {
+                                "CompressionType": "None",
+                                "S3DataType": "S3Prefix",
+                                "S3Uri": "s3://bucket/path/2",
+                            },
+                        },
+                    ]
+                },
+            },
+        },
     },
     "meta-textgeneration-llama-2-70b": {
         "Url": "https://ai.meta.com/resources/models-and-libraries/llama-downloads/",
diff --git a/tests/unit/sagemaker/jumpstart/hub/test_interfaces.py b/tests/unit/sagemaker/jumpstart/hub/test_interfaces.py
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py
