feat: Added utils for extracting JS data sources (#1471)

* added utils for accessing hosting data sources

* added utils for accessing hosting data sources

* removed other changes

* fixed formatting issues

* remove .keys()

* updated JumpStartModelDataSource

* fix slots

* remove print

* fix tests

* update tests

Author: akozd

src/sagemaker/jumpstart/types.py

@@ -884,41 +884,19 @@ def to_json(self) -> Dict[str, Any]:
         return json_obj
 
 
-class JumpStartModelDataSource(JumpStartDataHolderType):
+class JumpStartModelDataSource(AdditionalModelDataSource):
     """Data class JumpStart additional model data source."""
 
-    __slots__ = ["version", "additional_model_data_source"]
-
-    def __init__(self, spec: Dict[str, Any]):
-        """Initializes a JumpStartModelDataSource object.
-
-        Args:
-            spec (Dict[str, Any]): Dictionary representation of data source.
-        """
-        self.from_json(spec)
+    __slots__ = ["artifact_version"] + AdditionalModelDataSource.__slots__
 
     def from_json(self, json_obj: Dict[str, Any]) -> None:
         """Sets fields in object based on json.
 
         Args:
             json_obj (Dict[str, Any]): Dictionary representation of data source.
         """
-        self.version: str = json_obj["artifact_version"]
-        self.additional_model_data_source: AdditionalModelDataSource = AdditionalModelDataSource(
-            json_obj
-        )
-
-    def to_json(self) -> Dict[str, Any]:
-        """Returns json representation of JumpStartModelDataSource object."""
-        json_obj = {}
-        for att in self.__slots__:
-            if hasattr(self, att):
-                cur_val = getattr(self, att)
-                if issubclass(type(cur_val), JumpStartDataHolderType):
-                    json_obj[att] = cur_val.to_json()
-                else:
-                    json_obj[att] = cur_val
-        return json_obj
+        super().from_json(json_obj)
+        self.artifact_version: str = json_obj["artifact_version"]
 
 
 class JumpStartAdditionalDataSources(JumpStartDataHolderType):
@@ -1655,6 +1633,19 @@ def supports_incremental_training(self) -> bool:
         """Returns True if the model supports incremental training."""
         return self.incremental_training_supported
 
+    def get_speculative_decoding_s3_data_sources(self) -> List[JumpStartAdditionalDataSources]:
+        """Returns data sources for speculative decoding."""
+        return self.hosting_additional_data_sources.speculative_decoding or []
+
+    def get_additional_s3_data_sources(self) -> List[JumpStartAdditionalDataSources]:
+        """Returns a list of the additional S3 data sources for use by the model."""
+        additional_data_sources = []
+        if self.hosting_additional_data_sources:
+            for data_source in self.hosting_additional_data_sources.to_json():
+                data_sources = getattr(self.hosting_additional_data_sources, data_source) or []
+                additional_data_sources.extend(data_sources)
+        return additional_data_sources
+
 
 class JumpStartVersionedModelId(JumpStartDataHolderType):
     """Data class for versioned model IDs."""

tests/unit/sagemaker/jumpstart/constants.py

@@ -7518,6 +7518,37 @@
     "hosting_additional_data_sources": None,
 }
 
+BASE_HOSTING_ADDITIONAL_DATA_SOURCES = {
+    "hosting_additional_data_sources": {
+        "speculative_decoding": [
+            {
+                "channel_name": "speculative_decoding_channel",
+                "artifact_version": "version",
+                "s3_data_source": {
+                    "compression_type": "None",
+                    "s3_data_type": "S3Prefix",
+                    "s3_uri": "s3://bucket/path1",
+                    "hub_access_config": None,
+                    "model_access_config": None,
+                },
+            }
+        ],
+        "scripts": [
+            {
+                "channel_name": "scripts_channel",
+                "artifact_version": "version",
+                "s3_data_source": {
+                    "compression_type": "None",
+                    "s3_data_type": "S3Prefix",
+                    "s3_uri": "s3://bucket/path1",
+                    "hub_access_config": None,
+                    "model_access_config": None,
+                },
+            }
+        ],
+    },
+}
+
 BASE_HEADER = {
     "model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
     "version": "1.0.0",

tests/unit/sagemaker/jumpstart/test_types.py

@@ -28,6 +28,7 @@
 )
 from tests.unit.sagemaker.jumpstart.constants import (
     BASE_SPEC,
+    BASE_HOSTING_ADDITIONAL_DATA_SOURCES,
     INFERENCE_CONFIG_RANKINGS,
     INFERENCE_CONFIGS,
     TRAINING_CONFIG_RANKINGS,
@@ -436,6 +437,23 @@ def test_jumpstart_model_specs():
     assert specs3 == specs1
 
 
+def test_get_speculative_decoding_s3_data_sources():
+    specs = JumpStartModelSpecs({**BASE_SPEC, **BASE_HOSTING_ADDITIONAL_DATA_SOURCES})
+    assert (
+        specs.get_speculative_decoding_s3_data_sources()
+        == specs.hosting_additional_data_sources.speculative_decoding
+    )
+
+
+def test_get_additional_s3_data_sources():
+    specs = JumpStartModelSpecs({**BASE_SPEC, **BASE_HOSTING_ADDITIONAL_DATA_SOURCES})
+    data_sources = [
+        *specs.hosting_additional_data_sources.speculative_decoding,
+        *specs.hosting_additional_data_sources.scripts,
+    ]
+    assert specs.get_additional_s3_data_sources() == data_sources
+
+
 def test_jumpstart_image_uri_instance_variants():
 
     assert (