fix: linting, mypy, logical issues for jumpstart models

Author: evakravi

src/sagemaker/image_uris.py

@@ -92,7 +92,7 @@ def retrieve(
         )
         if image_scope is None:
             raise ValueError(
-                "Must specify `image_scope` argument to retrieve image uri for " "JumpStart models."
+                "Must specify `image_scope` argument to retrieve image uri for JumpStart models."
             )
         if image_scope == "inference":
             ecr_specs = model_specs.hosting_ecr_specs
@@ -103,7 +103,7 @@ def retrieve(
         else:
             raise ValueError("JumpStart models only support inference and training.")
 
-        if framework != None and framework != ecr_specs.framework:
+        if framework is not None and framework != ecr_specs.framework:
             raise ValueError(
                 f"Bad value for container framework for JumpStart model: '{framework}'."
             )

src/sagemaker/jumpstart/accessors.py

@@ -40,41 +40,78 @@ class JumpStartModelsCache(object):
     _cache: Optional[cache.JumpStartModelsCache] = None
     _curr_region = JUMPSTART_DEFAULT_REGION_NAME
 
-    _cache_kwargs = {}
+    _cache_kwargs: Dict[str, Any] = {}
 
+    @staticmethod
     def _validate_region_cache_kwargs(
-        cache_kwargs: Dict[str, Any] = {}, region: Optional[str] = None
-    ):
-        if region is not None and "region" in cache_kwargs:
-            if region != cache_kwargs["region"]:
+        cache_kwargs: Optional[Dict[str, Any]] = None, region: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Returns cache_kwargs with region argument removed if present.
+
+        Raises:
+            ValueError: If region in `cache_kwargs` is inconsistent with `region` argument.
+
+        Args:
+            cache_kwargs (Optional[Dict[str, Any]]): cache kwargs to validate.
+            region (str): The region to validate along with the kwargs.
+        """
+        cache_kwargs_dict = {} if cache_kwargs is None else cache_kwargs
+        assert isinstance(cache_kwargs_dict, dict)
+        if region is not None and "region" in cache_kwargs_dict:
+            if region != cache_kwargs_dict["region"]:
                 raise ValueError(
-                    f"Inconsistent region definitions: {region}, {cache_kwargs['region']}"
+                    f"Inconsistent region definitions: {region}, {cache_kwargs_dict['region']}"
                 )
-            del cache_kwargs["region"]
-        return cache_kwargs
+            del cache_kwargs_dict["region"]
+        return cache_kwargs_dict
 
     @staticmethod
     def get_model_header(region: str, model_id: str, version: str) -> JumpStartModelHeader:
+        """Returns model header from JumpStart models cache.
+
+        Args:
+            region (str): region for which to retrieve header.
+            model_id (str): model id to retrieve.
+            version (str): semantic version to retrieve for the model id.
+        """
         cache_kwargs = JumpStartModelsCache._validate_region_cache_kwargs(
             JumpStartModelsCache._cache_kwargs, region
         )
-        if JumpStartModelsCache._cache == None or region != JumpStartModelsCache._curr_region:
+        if JumpStartModelsCache._cache is None or region != JumpStartModelsCache._curr_region:
             JumpStartModelsCache._cache = cache.JumpStartModelsCache(region=region, **cache_kwargs)
             JumpStartModelsCache._curr_region = region
+        assert JumpStartModelsCache._cache is not None
         return JumpStartModelsCache._cache.get_header(model_id, version)
 
     @staticmethod
     def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelSpecs:
+        """Returns model specs from JumpStart models cache.
+
+        Args:
+            region (str): region for which to retrieve header.
+            model_id (str): model id to retrieve.
+            version (str): semantic version to retrieve for the model id.
+        """
         cache_kwargs = JumpStartModelsCache._validate_region_cache_kwargs(
             JumpStartModelsCache._cache_kwargs, region
         )
-        if JumpStartModelsCache._cache == None or region != JumpStartModelsCache._curr_region:
+        if JumpStartModelsCache._cache is None or region != JumpStartModelsCache._curr_region:
             JumpStartModelsCache._cache = cache.JumpStartModelsCache(region=region, **cache_kwargs)
             JumpStartModelsCache._curr_region = region
+        assert JumpStartModelsCache._cache is not None
         return JumpStartModelsCache._cache.get_specs(model_id, version)
 
     @staticmethod
     def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None:
+        """Sets cache kwargs. Clears the cache.
+
+        Raises:
+            ValueError: If region in `cache_kwargs` is inconsistent with `region` argument.
+
+        Args:
+            cache_kwargs (str): cache kwargs to validate.
+            region (str): The region to validate along with the kwargs.
+        """
         cache_kwargs = JumpStartModelsCache._validate_region_cache_kwargs(cache_kwargs, region)
         JumpStartModelsCache._cache_kwargs = cache_kwargs
         if region is None:
@@ -88,15 +125,15 @@ def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None:
             )
 
     @staticmethod
-    def reset_cache(cache_kwargs: Dict[str, Any] = {}, region: str = None) -> None:
-        cache_kwargs = JumpStartModelsCache._validate_region_cache_kwargs(cache_kwargs, region)
-        JumpStartModelsCache._cache_kwargs = cache_kwargs
-        if region is None:
-            JumpStartModelsCache._cache = cache.JumpStartModelsCache(
-                **JumpStartModelsCache._cache_kwargs
-            )
-        else:
-            JumpStartModelsCache._curr_region = region
-            JumpStartModelsCache._cache = cache.JumpStartModelsCache(
-                region=region, **JumpStartModelsCache._cache_kwargs
-            )
+    def reset_cache(cache_kwargs: Dict[str, Any] = None, region: str = None) -> None:
+        """Resets cache, optionally allowing cache kwargs to be passed to the new cache.
+
+        Raises:
+            ValueError: If region in `cache_kwargs` is inconsistent with `region` argument.
+
+        Args:
+            cache_kwargs (str): cache kwargs to validate.
+            region (str): The region to validate along with the kwargs.
+        """
+        cache_kwargs_dict = {} if cache_kwargs is None else cache_kwargs
+        JumpStartModelsCache.set_cache_kwargs(cache_kwargs_dict, region)

src/sagemaker/jumpstart/cache.py

@@ -48,37 +48,37 @@ class JumpStartModelsCache:
     for launching JumpStart models from the SageMaker SDK.
     """
 
+    # fmt: off
     def __init__(
         self,
-        region: Optional[str] = JUMPSTART_DEFAULT_REGION_NAME,
-        max_s3_cache_items: Optional[int] = JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
-        s3_cache_expiration_horizon: Optional[
-            datetime.timedelta
-        ] = JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
-        max_semantic_version_cache_items: Optional[
-            int
-        ] = JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
-        semantic_version_cache_expiration_horizon: Optional[
-            datetime.timedelta
-        ] = JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
-        manifest_file_s3_key: Optional[str] = JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
+        region: str = JUMPSTART_DEFAULT_REGION_NAME,
+        max_s3_cache_items: int = JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
+        s3_cache_expiration_horizon: datetime.timedelta =
+        JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
+        max_semantic_version_cache_items: int =
+        JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
+        semantic_version_cache_expiration_horizon: datetime.timedelta =
+        JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
+        manifest_file_s3_key: str =
+        JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
         s3_bucket_name: Optional[str] = None,
         s3_client_config: Optional[botocore.config.Config] = None,
-    ) -> None:
+    ) -> None:  # fmt: on
         """Initialize a ``JumpStartModelsCache`` instance.
 
         Args:
-            region (Optional[str]): AWS region to associate with cache. Default: region associated
+            region (str): AWS region to associate with cache. Default: region associated
                 with boto3 session.
-            max_s3_cache_items (Optional[int]): Maximum number of items to store in s3 cache.
+            max_s3_cache_items (int): Maximum number of items to store in s3 cache.
                 Default: 20.
-            s3_cache_expiration_horizon (Optional[datetime.timedelta]): Maximum time to hold
+            s3_cache_expiration_horizon (datetime.timedelta): Maximum time to hold
                 items in s3 cache before invalidation. Default: 6 hours.
-            max_semantic_version_cache_items (Optional[int]): Maximum number of items to store in
+            max_semantic_version_cache_items (int): Maximum number of items to store in
                 semantic version cache. Default: 20.
-            semantic_version_cache_expiration_horizon (Optional[datetime.timedelta]):
+            semantic_version_cache_expiration_horizon (datetime.timedelta):
                 Maximum time to hold items in semantic version cache before invalidation.
                 Default: 6 hours.
+            manifest_file_s3_key (str): The key in S3 corresponding to the sdk metadata manifest.
             s3_bucket_name (Optional[str]): S3 bucket to associate with cache.
                 Default: JumpStart-hosted content bucket for region.
             s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
@@ -126,7 +126,7 @@ def set_manifest_file_s3_key(self, key: str) -> None:
             self._manifest_file_s3_key = key
             self.clear()
 
-    def get_manifest_file_s3_key(self) -> None:
+    def get_manifest_file_s3_key(self) -> str:
         """Return manifest file s3 key for cache."""
         return self._manifest_file_s3_key
 
@@ -136,7 +136,7 @@ def set_s3_bucket_name(self, s3_bucket_name: str) -> None:
             self.s3_bucket_name = s3_bucket_name
             self.clear()
 
-    def get_bucket(self) -> None:
+    def get_bucket(self) -> str:
         """Return bucket used for cache."""
         return self.s3_bucket_name
 
@@ -166,6 +166,7 @@ def _get_manifest_key_from_model_id_semantic_version(
         manifest = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
         ).formatted_content
+        assert isinstance(manifest, dict)
 
         sm_version = utils.get_sagemaker_version()
 
@@ -191,16 +192,16 @@ def _get_manifest_key_from_model_id_semantic_version(
 
         if sm_incompatible_model_version is not None:
             model_version_to_use_incompatible_with_sagemaker = sm_incompatible_model_version
-            sm_version_to_use = [
+            sm_version_to_use_list = [
                 header.min_version
                 for header in manifest.values()
                 if header.model_id == model_id
                 and header.version == model_version_to_use_incompatible_with_sagemaker
             ]
-            if len(sm_version_to_use) != 1:
+            if len(sm_version_to_use_list) != 1:
                 # ``manifest`` dict should already enforce this
                 raise RuntimeError("Found more than one incompatible SageMaker version to use.")
-            sm_version_to_use = sm_version_to_use[0]
+            sm_version_to_use = sm_version_to_use_list[0]
 
             error_msg = (
                 f"Unable to find model manifest for {model_id} with version {version} "
@@ -258,9 +259,12 @@ def _get_file_from_s3(
     def get_manifest(self) -> List[JumpStartModelHeader]:
         """Return entire JumpStart models manifest."""
 
-        return self._s3_cache.get(
+        manifest_dict = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_content.values()
+        ).formatted_content
+        assert isinstance(manifest_dict, dict)
+        manifest = list(manifest_dict.values())
+        return manifest
 
     def get_header(self, model_id: str, semantic_version_str: str) -> JumpStartModelHeader:
         """Return header for a given JumpStart model id and semantic version.
@@ -277,30 +281,30 @@ def _select_version(
         self,
         semantic_version_str: str,
         available_versions: List[Version],
-    ) -> Optional[Version]:
-        """Utility to select appropriate version from available version given
-        a semantic version with which to filter.
+    ) -> Optional[str]:
+        """Utility to select appropriate version from available versions.
 
         Args:
             semantic_version_str (str): the semantic version for which to filter
                 available versions.
             available_versions (List[Version]): list of available versions.
         """
         if semantic_version_str == "*":
-            if len(available_versions) is 0:
+            if len(available_versions) == 0:
                 return None
-            else:
-                return str(max(available_versions))
-        else:
-            spec = SpecifierSet(f"=={semantic_version_str}")
-            available_versions = list(spec.filter(available_versions))
-            return str(available_versions[0]) if available_versions != [] else None
+            return str(max(available_versions))
+
+        spec = SpecifierSet(f"=={semantic_version_str}")
+        available_versions_filtered = list(spec.filter(available_versions))
+        return (
+            str(available_versions_filtered[0]) if available_versions_filtered != [] else None
+        )
 
     def _get_header_impl(
         self,
         model_id: str,
         semantic_version_str: str,
-        attempt: Optional[int] = 0,
+        attempt: int = 0,
     ) -> JumpStartModelHeader:
         """Lower-level function to return header.
 
@@ -310,7 +314,7 @@ def _get_header_impl(
             model_id (str): model id for which to get a header.
             semantic_version_str (str): The semantic version for which to get a
                 header.
-            attempt (Optional[int]): attempt number at retrieving a header.
+            attempt (int): attempt number at retrieving a header.
         """
 
         versioned_model_id = self._model_id_semantic_version_manifest_key_cache.get(
@@ -320,7 +324,10 @@ def _get_header_impl(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
         ).formatted_content
         try:
-            return manifest[versioned_model_id]
+            assert isinstance(manifest, dict)
+            header = manifest[versioned_model_id]
+            assert isinstance(header, JumpStartModelHeader)
+            return header
         except KeyError:
             if attempt > 0:
                 raise
@@ -338,9 +345,11 @@ def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelS
 
         header = self.get_header(model_id, semantic_version_str)
         spec_key = header.spec_key
-        return self._s3_cache.get(
+        specs = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.SPECS, spec_key)
         ).formatted_content
+        assert isinstance(specs, JumpStartModelSpecs)
+        return specs
 
     def clear(self) -> None:
         """Clears the model id/version and s3 cache."""

src/sagemaker/jumpstart/types.py

@@ -274,15 +274,15 @@ def __init__(
         self,
         formatted_content: Union[
             Dict[JumpStartVersionedModelId, JumpStartModelHeader],
-            List[JumpStartModelSpecs],
+            JumpStartModelSpecs,
         ],
         md5_hash: Optional[str] = None,
     ) -> None:
         """Instantiates JumpStartCachedS3ContentValue object.
 
         Args:
             formatted_content (Union[Dict[JumpStartVersionedModelId, JumpStartModelHeader],
-            List[JumpStartModelSpecs]]):
+            JumpStartModelSpecs]):
                 Formatted content for model specs and mappings from
                 versioned model ids to specs.
             md5_hash (str): md5_hash for stored file content from s3.

src/sagemaker/model_uris.py

@@ -13,10 +13,7 @@
 """Functions for generating S3 model artifact URIs for pre-built SageMaker models."""
 from __future__ import absolute_import
 
-import json
 import logging
-import os
-import re
 from typing import Optional
 
 from sagemaker.jumpstart import utils as jumpstart_utils
@@ -56,13 +53,15 @@ def retrieve(
     )
     if model_scope is None:
         raise ValueError(
-            "Must specify `model_scope` argument to retrieve model artifact uri for JumpStart models."
+            "Must specify `model_scope` argument to retrieve model "
+            "artifact uri for JumpStart models."
         )
     if model_scope == "inference":
         model_artifact_key = model_specs.hosting_artifact_key
     elif model_scope == "training":
         if not model_specs.training_supported:
             raise ValueError(f"JumpStart model id '{model_id}' does not support training.")
+        assert model_specs.training_artifact_key is not None
         model_artifact_key = model_specs.training_artifact_key
     else:
         raise ValueError("JumpStart models only support inference and training.")

src/sagemaker/script_uris.py

@@ -13,10 +13,7 @@
 """Functions for generating S3 model script URIs for pre-built SageMaker models."""
 from __future__ import absolute_import
 
-import json
 import logging
-import os
-import re
 
 from sagemaker.jumpstart import utils as jumpstart_utils
 from sagemaker.jumpstart import accessors as jumpstart_accessors