change: remove cache memory test, improve jumpstart cache design/tests

Author: evakravi

src/sagemaker/jumpstart/cache.py

@@ -234,11 +234,13 @@ def _get_file_from_s3(
         file_type, s3_key = key.file_type, key.s3_key
 
         if file_type == JumpStartS3FileType.MANIFEST:
-            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
+            if value is not None:
+                etag = self._s3_client.head_object(Bucket=self.s3_bucket_name, Key=s3_key)["ETag"]
+                if etag == value.md5_hash:
+                    return value
             response = self._s3_client.get_object(Bucket=self.s3_bucket_name, Key=s3_key)
             formatted_body = json.loads(response["Body"].read().decode("utf-8"))
+            etag = response["ETag"]
             return JumpStartCachedS3ContentValue(
                 formatted_file_content=utils.get_formatted_manifest(formatted_body),
                 md5_hash=etag,
@@ -271,10 +273,10 @@ def get_header(
                 header. If None, the highest compatible version is returned.
         """
 
-        return self._get_header_impl(model_id, 0, semantic_version_str)
+        return self._get_header_impl(model_id, semantic_version_str=semantic_version_str)
 
     def _get_header_impl(
-        self, model_id: str, attempt: int, semantic_version_str: Optional[str] = None
+        self, model_id: str, attempt: Optional[int] = 0, semantic_version_str: Optional[str] = None
     ) -> JumpStartModelHeader:
         """Lower-level function to return header.
 

src/sagemaker/jumpstart/types.py

@@ -144,12 +144,12 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "min_sdk_version",
         "incremental_training_supported",
         "hosting_ecr_specs",
-        "hosting_artifact_uri",
-        "hosting_script_uri",
+        "hosting_artifact_key",
+        "hosting_script_key",
         "training_supported",
         "training_ecr_specs",
-        "training_artifact_uri",
-        "training_script_uri",
+        "training_artifact_key",
+        "training_script_key",
         "hyperparameters",
     ]
 
@@ -164,20 +164,20 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.min_sdk_version: str = json_obj["min_sdk_version"]
         self.incremental_training_supported: bool = bool(json_obj["incremental_training_supported"])
         self.hosting_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(json_obj["hosting_ecr_specs"])
-        self.hosting_artifact_uri: str = json_obj["hosting_artifact_uri"]
-        self.hosting_script_uri: str = json_obj["hosting_script_uri"]
+        self.hosting_artifact_key: str = json_obj["hosting_artifact_key"]
+        self.hosting_script_key: str = json_obj["hosting_script_key"]
         self.training_supported: bool = bool(json_obj["training_supported"])
         if self.training_supported:
             self.training_ecr_specs: Optional[JumpStartECRSpecs] = JumpStartECRSpecs(
                 json_obj["training_ecr_specs"]
             )
-            self.training_artifact_uri: Optional[str] = json_obj["training_artifact_uri"]
-            self.training_script_uri: Optional[str] = json_obj["training_script_uri"]
+            self.training_artifact_key: Optional[str] = json_obj["training_artifact_key"]
+            self.training_script_key: Optional[str] = json_obj["training_script_key"]
             self.hyperparameters: Optional[Dict[str, Any]] = json_obj["hyperparameters"]
         else:
             self.training_ecr_specs = (
-                self.training_artifact_uri
-            ) = self.training_script_uri = self.hyperparameters = None
+                self.training_artifact_key
+            ) = self.training_script_key = self.hyperparameters = None
 
     def to_json(self) -> Dict[str, Any]:
         """Returns json representation of JumpStartModelSpecs object."""

src/sagemaker/jumpstart/utils.py

@@ -82,7 +82,8 @@ def get_sagemaker_version() -> str:
     rejected.
 
     Raises:
-        RuntimeError: If the SageMaker version is not readable.
+        RuntimeError: If the SageMaker version is not readable. An exception is also raised if
+        the version cannot be parsed by ``semantic_version``.
     """
     version = sagemaker.__version__
     parsed_version = None

src/sagemaker/utilities/cache.py

@@ -134,12 +134,12 @@ def put(self, key: KeyType, value: Optional[ValType] = None) -> None:
     def _get_item(self, key: KeyType, fail_on_old_value: bool) -> ValType:
         """Returns value from cache corresponding to key.
 
-        If ``fail_on_old_value``, a KeyError is thrown instead of a new value
+        If ``fail_on_old_value``, a KeyError is raised instead of a new value
             getting fetched.
 
         Args:
             key (KeyType): Key in cache to retrieve.
-            fail_on_old_value (bool): True if a KeyError is thrown when the cache value
+            fail_on_old_value (bool): True if a KeyError is raised when the cache value
                 is old.
 
         Raises:

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -49,10 +49,10 @@
         "framework_version": "1.9.0",
         "py_version": "py3",
     },
-    "hosting_artifact_uri": "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz",
-    "training_artifact_uri": "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz",
-    "hosting_script_uri": "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
-    "training_script_uri": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
+    "hosting_artifact_key": "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz",
+    "training_artifact_key": "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz",
+    "hosting_script_key": "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
+    "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
     "hyperparameters": {
         "adam-learning-rate": {"type": "float", "default": 0.05, "min": 1e-08, "max": 1},
         "epochs": {"type": "int", "default": 3, "min": 1, "max": 1000},
@@ -334,30 +334,100 @@ def test_jumpstart_cache_gets_cleared_when_params_are_set(mock_boto3_client):
 
 
 def test_jumpstart_cache_handles_boto3_client_errors():
+    # Testing get_object
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
     stubbed_s3_client = Stubber(cache._s3_client)
-    stubbed_s3_client.add_client_error("head_object", http_status_code=404)
     stubbed_s3_client.add_client_error("get_object", http_status_code=404)
     stubbed_s3_client.activate()
     with pytest.raises(botocore.exceptions.ClientError):
         cache.get_header(model_id="tensorflow-ic-imagenet-inception-v3-classification-4")
 
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
     stubbed_s3_client = Stubber(cache._s3_client)
-    stubbed_s3_client.add_client_error("head_object", service_error_code="AccessDenied")
     stubbed_s3_client.add_client_error("get_object", service_error_code="AccessDenied")
     stubbed_s3_client.activate()
     with pytest.raises(botocore.exceptions.ClientError):
         cache.get_header(model_id="tensorflow-ic-imagenet-inception-v3-classification-4")
 
     cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
     stubbed_s3_client = Stubber(cache._s3_client)
-    stubbed_s3_client.add_client_error("head_object", service_error_code="EndpointConnectionError")
     stubbed_s3_client.add_client_error("get_object", service_error_code="EndpointConnectionError")
     stubbed_s3_client.activate()
     with pytest.raises(botocore.exceptions.ClientError):
         cache.get_header(model_id="tensorflow-ic-imagenet-inception-v3-classification-4")
 
+    # Testing head_object:
+    mock_now = datetime.datetime.fromtimestamp(1636730651.079551)
+    with patch("datetime.datetime") as mock_datetime:
+        mock_manifest_json = json.dumps(
+            [
+                {
+                    "model_id": "pytorch-ic-imagenet-inception-v3-classification-4",
+                    "version": "2.0.0",
+                    "min_version": "2.49.0",
+                    "spec_key": "community_models_specs/pytorch-ic-"
+                    "imagenet-inception-v3-classification-4/specs_v2.0.0.json",
+                }
+            ]
+        )
+
+        get_object_mocked_response = {
+            "Body": botocore.response.StreamingBody(
+                io.BytesIO(bytes(mock_manifest_json, "utf-8")),
+                content_length=len(mock_manifest_json),
+            ),
+            "ETag": "etag",
+        }
+
+        mock_datetime.now.return_value = mock_now
+
+        cache1 = JumpStartModelsCache(
+            s3_bucket_name="some_bucket", s3_cache_expiration_horizon=datetime.timedelta(hours=1)
+        )
+        stubbed_s3_client1 = Stubber(cache1._s3_client)
+
+        stubbed_s3_client1.add_response("get_object", copy.deepcopy(get_object_mocked_response))
+        stubbed_s3_client1.activate()
+        cache1.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
+
+        mock_datetime.now.return_value += datetime.timedelta(weeks=1)
+
+        stubbed_s3_client1.add_client_error("head_object", http_status_code=404)
+        with pytest.raises(botocore.exceptions.ClientError):
+            cache1.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
+
+        cache2 = JumpStartModelsCache(
+            s3_bucket_name="some_bucket", s3_cache_expiration_horizon=datetime.timedelta(hours=1)
+        )
+        stubbed_s3_client2 = Stubber(cache2._s3_client)
+
+        stubbed_s3_client2.add_response("get_object", copy.deepcopy(get_object_mocked_response))
+        stubbed_s3_client2.activate()
+        cache2.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
+
+        mock_datetime.now.return_value += datetime.timedelta(weeks=1)
+
+        stubbed_s3_client2.add_client_error("head_object", service_error_code="AccessDenied")
+        with pytest.raises(botocore.exceptions.ClientError):
+            cache2.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
+
+        cache3 = JumpStartModelsCache(
+            s3_bucket_name="some_bucket", s3_cache_expiration_horizon=datetime.timedelta(hours=1)
+        )
+        stubbed_s3_client3 = Stubber(cache3._s3_client)
+
+        stubbed_s3_client3.add_response("get_object", copy.deepcopy(get_object_mocked_response))
+        stubbed_s3_client3.activate()
+        cache3.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
+
+        mock_datetime.now.return_value += datetime.timedelta(weeks=1)
+
+        stubbed_s3_client3.add_client_error(
+            "head_object", service_error_code="EndpointConnectionError"
+        )
+        with pytest.raises(botocore.exceptions.ClientError):
+            cache3.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
+
 
 def test_jumpstart_cache_accepts_input_parameters():
 
@@ -422,19 +492,18 @@ def test_jumpstart_cache_evaluates_md5_hash(mock_boto3_client):
         mock_boto3_client.return_value.get_object.return_value = {
             "Body": botocore.response.StreamingBody(
                 io.BytesIO(bytes(mock_json, "utf-8")), content_length=len(mock_json)
-            )
+            ),
+            "ETag": "hash1",
         }
         mock_boto3_client.return_value.head_object.return_value = {"ETag": "hash1"}
 
         cache.get_header(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
 
-        # first time accessing cache should involve get_object and head_object
+        # first time accessing cache should just involve get_object
         mock_boto3_client.return_value.get_object.assert_called_with(
             Bucket=bucket_name, Key=JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY
         )
-        mock_boto3_client.return_value.head_object.assert_called_with(
-            Bucket=bucket_name, Key=JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY
-        )
+        mock_boto3_client.return_value.head_object.assert_not_called()
 
         mock_boto3_client.return_value.get_object.reset_mock()
         mock_boto3_client.return_value.head_object.reset_mock()
@@ -443,7 +512,8 @@ def test_jumpstart_cache_evaluates_md5_hash(mock_boto3_client):
         mock_boto3_client.return_value.get_object.return_value = {
             "Body": botocore.response.StreamingBody(
                 io.BytesIO(bytes(mock_json, "utf-8")), content_length=len(mock_json)
-            )
+            ),
+            "ETag": "hash1",
         }
         mock_boto3_client.return_value.head_object.return_value = {"ETag": "hash1"}
 
@@ -465,7 +535,8 @@ def test_jumpstart_cache_evaluates_md5_hash(mock_boto3_client):
         mock_boto3_client.return_value.get_object.return_value = {
             "Body": botocore.response.StreamingBody(
                 io.BytesIO(bytes(mock_json, "utf-8")), content_length=len(mock_json)
-            )
+            ),
+            "ETag": "hash2",
         }
 
         # invalidate cache
@@ -499,7 +570,8 @@ def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
     mock_boto3_client.return_value.get_object.return_value = {
         "Body": botocore.response.StreamingBody(
             io.BytesIO(bytes(mock_json, "utf-8")), content_length=len(mock_json)
-        )
+        ),
+        "ETag": "etag",
     }
 
     mock_boto3_client.return_value.head_object.return_value = {"ETag": "some-hash"}
@@ -514,9 +586,8 @@ def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
     mock_boto3_client.return_value.get_object.assert_called_with(
         Bucket=bucket_name, Key=JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY
     )
-    mock_boto3_client.return_value.head_object.assert_called_with(
-        Bucket=bucket_name, Key=JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY
-    )
+    mock_boto3_client.return_value.head_object.assert_not_called()
+
     mock_boto3_client.assert_called_with("s3", region_name="my_region", config=client_config)
 
     # test get_specs. manifest already in cache, so only s3 call will be to get specs.
@@ -527,7 +598,8 @@ def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
     mock_boto3_client.return_value.get_object.return_value = {
         "Body": botocore.response.StreamingBody(
             io.BytesIO(bytes(mock_json, "utf-8")), content_length=len(mock_json)
-        )
+        ),
+        "ETag": "etag",
     }
     cache.get_specs(model_id="pytorch-ic-imagenet-inception-v3-classification-4")
 

tests/unit/sagemaker/jumpstart/test_types.py

@@ -69,10 +69,10 @@ def test_jumpstart_model_specs():
             "framework_version": "1.9.0",
             "py_version": "py3",
         },
-        "hosting_artifact_uri": "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz",
-        "training_artifact_uri": "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz",
-        "hosting_script_uri": "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
-        "training_script_uri": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
+        "hosting_artifact_key": "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz",
+        "training_artifact_key": "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz",
+        "hosting_script_key": "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
+        "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
         "hyperparameters": {
             "adam-learning-rate": {"type": "float", "default": 0.05, "min": 1e-08, "max": 1},
             "epochs": {"type": "int", "default": 3, "min": 1, "max": 1000},
@@ -101,14 +101,14 @@ def test_jumpstart_model_specs():
             "py_version": "py3",
         }
     )
-    assert specs1.hosting_artifact_uri == "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz"
-    assert specs1.training_artifact_uri == "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz"
+    assert specs1.hosting_artifact_key == "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz"
+    assert specs1.training_artifact_key == "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz"
     assert (
-        specs1.hosting_script_uri
+        specs1.hosting_script_key
         == "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz"
     )
     assert (
-        specs1.training_script_uri
+        specs1.training_script_key
         == "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz"
     )
     assert specs1.hyperparameters == {

tests/unit/sagemaker/utilities/test_cache.py

@@ -14,7 +14,6 @@
 from typing import Optional, Union
 from mock.mock import MagicMock, patch
 import pytest
-import pickle
 
 
 from sagemaker.utilities import cache
@@ -54,16 +53,16 @@ def test_cache_invalidates_old_item():
         retrieval_function=retrieval_function,
     )
 
-    curr_time = datetime.datetime.fromtimestamp(1636730651.079551)
+    mock_curr_time = datetime.datetime.fromtimestamp(1636730651.079551)
     with patch("datetime.datetime") as mock_datetime:
-        mock_datetime.now.return_value = curr_time
+        mock_datetime.now.return_value = mock_curr_time
         my_cache.put(5)
         mock_datetime.now.return_value += datetime.timedelta(milliseconds=2)
         with pytest.raises(KeyError):
             my_cache.get(5, False)
 
     with patch("datetime.datetime") as mock_datetime:
-        mock_datetime.now.return_value = curr_time
+        mock_datetime.now.return_value = mock_curr_time
         my_cache.put(5)
         mock_datetime.now.return_value += datetime.timedelta(milliseconds=0.5)
         assert my_cache.get(5, False) == retrieval_function(key=5)
@@ -76,15 +75,15 @@ def test_cache_fetches_new_item():
         retrieval_function=retrieval_function,
     )
 
-    curr_time = datetime.datetime.fromtimestamp(1636730651.079551)
+    mock_curr_time = datetime.datetime.fromtimestamp(1636730651.079551)
     with patch("datetime.datetime") as mock_datetime:
-        mock_datetime.now.return_value = curr_time
+        mock_datetime.now.return_value = mock_curr_time
         my_cache.put(5, 10)
         mock_datetime.now.return_value += datetime.timedelta(milliseconds=2)
         assert my_cache.get(5) == retrieval_function(key=5)
 
     with patch("datetime.datetime") as mock_datetime:
-        mock_datetime.now.return_value = curr_time
+        mock_datetime.now.return_value = mock_curr_time
         my_cache.put(5, 10)
         mock_datetime.now.return_value += datetime.timedelta(milliseconds=0.5)
         assert my_cache.get(5, False) == 10
@@ -194,24 +193,3 @@ def test_cache_clear_and_contains():
     assert len(my_cache) == 0
     with pytest.raises(KeyError):
         my_cache.get(1, False)
-
-
-def test_cache_memory_usage():
-    my_cache = cache.LRUCache[int, Union[int, str]](
-        max_cache_items=10,
-        expiration_horizon=datetime.timedelta(hours=1),
-        retrieval_function=retrieval_function,
-    )
-    cache_size_bytes = []
-    cache_size_bytes.append(len(pickle.dumps(my_cache)))
-    for i in range(50):
-        my_cache.put(i)
-        cache_size_bytes.append(len(pickle.dumps(my_cache)))
-
-    max_cache_items_iter_cache_size = cache_size_bytes[10]
-    past_capacity_iter_cache_size = cache_size_bytes[50]
-    percent_difference_cache_size = (
-        abs(past_capacity_iter_cache_size - max_cache_items_iter_cache_size)
-        / max_cache_items_iter_cache_size
-    ) * 100.0
-    assert percent_difference_cache_size < 1