fix: update jumpstart cache and unit tests

Author: evakravi

src/sagemaker/jumpstart/cache.py

@@ -13,10 +13,21 @@
 """This module defines the JumpStartModelsCache class."""
 from __future__ import absolute_import
 import datetime
-from typing import List, Optional
+from typing import Optional
 import json
 import boto3
+import botocore
 import semantic_version
+from sagemaker.jumpstart.constants import (
+    JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
+from sagemaker.jumpstart.parameters import (
+    JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
+    JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
+    JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
+    JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
+)
 from sagemaker.jumpstart.types import (
     JumpStartCachedS3ContentKey,
     JumpStartCachedS3ContentValue,
@@ -28,16 +39,6 @@
 from sagemaker.jumpstart import utils
 from sagemaker.utilities.cache import LRUCache
 
-DEFAULT_REGION_NAME = boto3.session.Session().region_name
-
-DEFAULT_MAX_S3_CACHE_ITEMS = 20
-DEFAULT_S3_CACHE_EXPIRATION_TIME = datetime.timedelta(hours=6)
-
-DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS = 20
-DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_TIME = datetime.timedelta(hours=6)
-
-DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
-
 
 class JumpStartModelsCache:
     """Class that implements a cache for JumpStart models manifests and specs.
@@ -48,78 +49,95 @@ class JumpStartModelsCache:
 
     def __init__(
         self,
-        region: Optional[str] = DEFAULT_REGION_NAME,
-        max_s3_cache_items: Optional[int] = DEFAULT_MAX_S3_CACHE_ITEMS,
-        s3_cache_expiration_time: Optional[datetime.timedelta] = DEFAULT_S3_CACHE_EXPIRATION_TIME,
-        max_semantic_version_cache_items: Optional[int] = DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
-        semantic_version_cache_expiration_time: Optional[
+        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
-        ] = DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_TIME,
-        manifest_file_s3_key: Optional[str] = DEFAULT_MANIFEST_FILE_S3_KEY,
-        bucket: Optional[str] = None,
+        ] = 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,
+        s3_bucket_name: Optional[str] = None,
+        s3_client_config: Optional[botocore.config.Config] = None,
     ) -> None:
         """Initialize a ``JumpStartModelsCache`` instance.
 
         Args:
             region (Optional[str]): AWS region to associate with cache. Default: region associated
-                with botocore session.
-            max_s3_cache_items (Optional[int]): Maximum number of files to store in s3 cache.
+                with boto3 session.
+            max_s3_cache_items (Optional[int]): Maximum number of items to store in s3 cache.
                 Default: 20.
-            s3_cache_expiration_time (Optional[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 files to store in
+            s3_cache_expiration_horizon (Optional[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
                 semantic version cache. Default: 20.
-            semantic_version_cache_expiration_time (Optional[datetime.timedelta]): Maximum time to
-                hold items in semantic version cache before invalidation. Default: 6 hours.
-            bucket (Optional[str]): S3 bucket to associate with cache. Default: JumpStart-hosted
-                content bucket for region.
+            semantic_version_cache_expiration_horizon (Optional[datetime.timedelta]):
+                Maximum time to hold items in semantic version cache before invalidation.
+                Default: 6 hours.
+            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.
+                Default: None (no config).
         """
 
         self._region = region
         self._s3_cache = LRUCache[JumpStartCachedS3ContentKey, JumpStartCachedS3ContentValue](
             max_cache_items=max_s3_cache_items,
-            expiration_time=s3_cache_expiration_time,
+            expiration_horizon=s3_cache_expiration_horizon,
             retrieval_function=self._get_file_from_s3,
         )
         self._model_id_semantic_version_manifest_key_cache = LRUCache[
             JumpStartVersionedModelId, JumpStartVersionedModelId
         ](
             max_cache_items=max_semantic_version_cache_items,
-            expiration_time=semantic_version_cache_expiration_time,
+            expiration_horizon=semantic_version_cache_expiration_horizon,
             retrieval_function=self._get_manifest_key_from_model_id_semantic_version,
         )
         self._manifest_file_s3_key = manifest_file_s3_key
-        self._bucket = (
-            utils.get_jumpstart_content_bucket(self._region) if bucket is None else bucket
+        self.s3_bucket_name = (
+            utils.get_jumpstart_content_bucket(self._region)
+            if s3_bucket_name is None
+            else s3_bucket_name
+        )
+        self._s3_client = (
+            boto3.client("s3", region_name=self._region, config=s3_client_config)
+            if s3_client_config
+            else boto3.client("s3", region_name=self._region)
         )
-        self._has_retried_cache_refresh = False
 
     def set_region(self, region: str) -> None:
         """Set region for cache. Clears cache after new region is set."""
-        self._region = region
-        self.clear()
+        if region != self._region:
+            self._region = region
+            self.clear()
 
     def get_region(self) -> str:
         """Return region for cache."""
         return self._region
 
     def set_manifest_file_s3_key(self, key: str) -> None:
         """Set manifest file s3 key. Clears cache after new key is set."""
-        self._manifest_file_s3_key = key
-        self.clear()
+        if key != self._manifest_file_s3_key:
+            self._manifest_file_s3_key = key
+            self.clear()
 
     def get_manifest_file_s3_key(self) -> None:
         """Return manifest file s3 key for cache."""
         return self._manifest_file_s3_key
 
-    def set_bucket(self, bucket: str) -> None:
+    def set_s3_bucket_name(self, s3_bucket_name: str) -> None:
         """Set s3 bucket used for cache."""
-        self._bucket = bucket
-        self.clear()
+        if s3_bucket_name != self.s3_bucket_name:
+            self.s3_bucket_name = s3_bucket_name
+            self.clear()
 
     def get_bucket(self) -> None:
         """Return bucket used for cache."""
-        return self._bucket
+        return self.s3_bucket_name
 
     def _get_manifest_key_from_model_id_semantic_version(
         self,
@@ -128,13 +146,18 @@ def _get_manifest_key_from_model_id_semantic_version(
     ) -> JumpStartVersionedModelId:
         """Return model id and version in manifest that matches semantic version/id.
 
+        Uses ``semantic_version`` to perform version comparison. The highest model version
+        matching the semantic version is used, which is compatible with the SageMaker
+        version.
+
         Args:
             key (JumpStartVersionedModelId): Key for which to fetch versioned model id.
             value (Optional[JumpStartVersionedModelId]): Unused variable for current value of
                 old cached model id/version.
 
         Raises:
-            KeyError: If the semantic version is not found in the manifest.
+            KeyError: If the semantic version is not found in the manifest, or is found but
+                the SageMaker version needs to be upgraded in order for the model to be used.
         """
 
         model_id, version = key.model_id, key.version
@@ -147,7 +170,7 @@ def _get_manifest_key_from_model_id_semantic_version(
 
         versions_compatible_with_sagemaker = [
             semantic_version.Version(header.version)
-            for _, header in manifest.items()
+            for header in manifest.values()
             if header.model_id == model_id
             and semantic_version.Version(header.min_version) <= semantic_version.Version(sm_version)
         ]
@@ -164,19 +187,19 @@ def _get_manifest_key_from_model_id_semantic_version(
 
         versions_incompatible_with_sagemaker = [
             semantic_version.Version(header.version)
-            for _, header in manifest.items()
+            for header in manifest.values()
             if header.model_id == model_id
         ]
         sm_incompatible_model_version = spec.select(versions_incompatible_with_sagemaker)
         if sm_incompatible_model_version is not None:
             model_version_to_use_incompatible_with_sagemaker = str(sm_incompatible_model_version)
             sm_version_to_use = [
                 header.min_version
-                for _, header in manifest.items()
+                for header in manifest.values()
                 if header.model_id == model_id
                 and header.version == model_version_to_use_incompatible_with_sagemaker
             ]
-            assert len(sm_version_to_use) == 1
+            assert len(sm_version_to_use) == 1  # ``manifest`` dict should already enforce this
             sm_version_to_use = sm_version_to_use[0]
 
             error_msg = (
@@ -187,7 +210,7 @@ def _get_manifest_key_from_model_id_semantic_version(
                 f"{model_version_to_use_incompatible_with_sagemaker} of {model_id}."
             )
             raise KeyError(error_msg)
-        error_msg = f"Unable to find model manifest for {model_id} with version {version}"
+        error_msg = f"Unable to find model manifest for {model_id} with version {version}."
         raise KeyError(error_msg)
 
     def _get_file_from_s3(
@@ -210,33 +233,49 @@ def _get_file_from_s3(
 
         file_type, s3_key = key.file_type, key.s3_key
 
-        s3_client = boto3.client("s3", region_name=self._region)
-
         if file_type == JumpStartS3FileType.MANIFEST:
-            etag = s3_client.head_object(Bucket=self._bucket, Key=s3_key)["ETag"]
+            etag = self._s3_client.head_object(Bucket=self.s3_bucket_name, Key=s3_key)["ETag"]
             if value is not None and etag == value.md5_hash:
                 return value
-            response = s3_client.get_object(Bucket=self._bucket, Key=s3_key)
+            response = self._s3_client.get_object(Bucket=self.s3_bucket_name, Key=s3_key)
             formatted_body = json.loads(response["Body"].read().decode("utf-8"))
             return JumpStartCachedS3ContentValue(
                 formatted_file_content=utils.get_formatted_manifest(formatted_body),
                 md5_hash=etag,
             )
         if file_type == JumpStartS3FileType.SPECS:
-            response = s3_client.get_object(Bucket=self._bucket, Key=s3_key)
+            response = self._s3_client.get_object(Bucket=self.s3_bucket_name, Key=s3_key)
             formatted_body = json.loads(response["Body"].read().decode("utf-8"))
             return JumpStartCachedS3ContentValue(
                 formatted_file_content=JumpStartModelSpecs(formatted_body)
             )
-        raise RuntimeError(f"Bad value for key: {key}")
+        raise ValueError(
+            f"Bad value for key '{key}': must be in {[JumpStartS3FileType.MANIFEST, JumpStartS3FileType.SPECS]}"
+        )
 
     def get_header(
         self, model_id: str, semantic_version_str: Optional[str] = None
-    ) -> List[JumpStartModelHeader]:
-        """Return list of headers for a given JumpStart model id and semantic version.
+    ) -> JumpStartModelHeader:
+        """Return header for a given JumpStart model id and semantic version.
+
+        Args:
+            model_id (str): model id for which to get a header.
+            semantic_version_str (Optional[str]): The semantic version for which to get a
+                header. If None, the highest compatible version is returned.
+        """
+
+        return self._get_header_impl(model_id, 0, semantic_version_str)
+
+    def _get_header_impl(
+        self, model_id: str, attempt: int, semantic_version_str: Optional[str] = None
+    ) -> JumpStartModelHeader:
+        """Lower-level function to return header.
+
+        Allows a single retry if the cache is old.
 
         Args:
             model_id (str): model id for which to get a header.
+            attempt (int): attempt number at retrieving a header.
             semantic_version_str (Optional[str]): The semantic version for which to get a
                 header. If None, the highest compatible version is returned.
         """
@@ -248,17 +287,12 @@ def get_header(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
         ).formatted_file_content
         try:
-            header = manifest[versioned_model_id]
-            if self._has_retried_cache_refresh:
-                self._has_retried_cache_refresh = False
-            return header
+            return manifest[versioned_model_id]
         except KeyError:
-            if self._has_retried_cache_refresh:
-                self._has_retried_cache_refresh = False
+            if attempt > 0:
                 raise
             self.clear()
-            self._has_retried_cache_refresh = True
-            return self.get_header(model_id, semantic_version)
+            return self._get_header_impl(model_id, attempt + 1, semantic_version_str)
 
     def get_specs(
         self, model_id: str, semantic_version_str: Optional[str] = None
@@ -278,7 +312,6 @@ def get_specs(
         ).formatted_file_content
 
     def clear(self) -> None:
-        """Clears the model id/version and s3 cache and resets ``_has_retried_cache_refresh``."""
+        """Clears the model id/version and s3 cache."""
         self._s3_cache.clear()
         self._model_id_semantic_version_manifest_key_cache.clear()
-        self._has_retried_cache_refresh = False

src/sagemaker/jumpstart/constants.py

@@ -13,10 +13,17 @@
 """This module stores constants related to SageMaker JumpStart."""
 from __future__ import absolute_import
 from typing import Set
+import boto3
 from sagemaker.jumpstart.types import JumpStartLaunchedRegionInfo
 
 
-LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set()
+JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set()
 
-REGION_NAME_TO_LAUNCHED_REGION_DICT = {region.region_name: region for region in LAUNCHED_REGIONS}
-REGION_NAME_SET = {region.region_name for region in LAUNCHED_REGIONS}
+JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT = {
+    region.region_name: region for region in JUMPSTART_LAUNCHED_REGIONS
+}
+JUMPSTART_REGION_NAME_SET = {region.region_name for region in JUMPSTART_LAUNCHED_REGIONS}
+
+JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name
+
+JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"

src/sagemaker/jumpstart/parameters.py

@@ -0,0 +1,20 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module stores parameters related to SageMaker JumpStart."""
+from __future__ import absolute_import
+import datetime
+
+JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS = 20
+JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS = 20
+JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON = datetime.timedelta(hours=6)
+JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON = datetime.timedelta(hours=6)