feat: add proprietary manifest/specs parsing

add unittests for test_cache

small refactoring

address comments and more unittests

fix linting and fix more tests

fix: pylint

feat: JumpStartModel class for prop models

Author: Captainia

src/sagemaker/accept_types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -72,6 +73,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,

src/sagemaker/base_predictor.py

@@ -58,7 +58,9 @@
 from sagemaker.model_monitor.model_monitoring import DEFAULT_REPOSITORY_NAME
 
 from sagemaker.lineage.context import EndpointContext
-from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
+from sagemaker.compute_resource_requirements.resource_requirements import (
+    ResourceRequirements,
+)
 
 LOGGER = logging.getLogger("sagemaker")
 

src/sagemaker/content_types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -72,6 +73,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,

src/sagemaker/deserializers.py

@@ -35,6 +35,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -95,6 +96,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
 ) -> BaseDeserializer:
     """Retrieves the default deserializer for the model matching the given arguments.
 

src/sagemaker/jumpstart/accessors.py

@@ -18,6 +18,7 @@
 
 from sagemaker.deprecations import deprecated
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.jumpstart import cache
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
@@ -197,7 +198,9 @@ def _set_cache_and_region(region: str, cache_kwargs: dict) -> None:
 
     @staticmethod
     def _get_manifest(
-        region: str = JUMPSTART_DEFAULT_REGION_NAME, s3_client: Optional[boto3.client] = None
+        region: str = JUMPSTART_DEFAULT_REGION_NAME,
+        s3_client: Optional[boto3.client] = None,
+        model_type: JumpStartModelType = JumpStartModelType.OPEN_SOURCE,
     ) -> List[JumpStartModelHeader]:
         """Return entire JumpStart models manifest.
 
@@ -215,13 +218,19 @@ def _get_manifest(
             additional_kwargs.update({"s3_client": s3_client})
 
         cache_kwargs = JumpStartModelsAccessor._validate_and_mutate_region_cache_kwargs(
-            {**JumpStartModelsAccessor._cache_kwargs, **additional_kwargs}, region
+            {**JumpStartModelsAccessor._cache_kwargs, **additional_kwargs},
+            region,
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
-        return JumpStartModelsAccessor._cache.get_manifest()  # type: ignore
+        return JumpStartModelsAccessor._cache.get_manifest(model_type)  # type: ignore
 
     @staticmethod
-    def get_model_header(region: str, model_id: str, version: str) -> JumpStartModelHeader:
+    def get_model_header(
+        region: str,
+        model_id: str,
+        version: str,
+        model_type: JumpStartModelType = JumpStartModelType.OPEN_SOURCE,
+    ) -> JumpStartModelHeader:
         """Returns model header from JumpStart models cache.
 
         Args:
@@ -234,12 +243,18 @@ def get_model_header(region: str, model_id: str, version: str) -> JumpStartModel
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
         return JumpStartModelsAccessor._cache.get_header(  # type: ignore
-            model_id=model_id, semantic_version_str=version
+            model_id=model_id,
+            semantic_version_str=version,
+            model_type=model_type,
         )
 
     @staticmethod
     def get_model_specs(
-        region: str, model_id: str, version: str, s3_client: Optional[boto3.client] = None
+        region: str,
+        model_id: str,
+        version: str,
+        s3_client: Optional[boto3.client] = None,
+        model_type=JumpStartModelType.OPEN_SOURCE,
     ) -> JumpStartModelSpecs:
         """Returns model specs from JumpStart models cache.
 
@@ -260,7 +275,7 @@ def get_model_specs(
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
-            model_id=model_id, semantic_version_str=version
+            model_id=model_id, version_str=version, model_type=model_type
         )
 
     @staticmethod