Merge branch 'master' into fix/tags-js-model-packages

Author: evakravi

src/sagemaker/jumpstart/accessors.py

@@ -42,8 +42,20 @@ class JumpStartModelsAccessor(object):
     _cache: Optional[cache.JumpStartModelsCache] = None
     _curr_region = JUMPSTART_DEFAULT_REGION_NAME
 
+    _content_bucket: Optional[str] = None
+
     _cache_kwargs: Dict[str, Any] = {}
 
+    @staticmethod
+    def set_jumpstart_content_bucket(content_bucket: str) -> None:
+        """Sets JumpStart content bucket."""
+        JumpStartModelsAccessor._content_bucket = content_bucket
+
+    @staticmethod
+    def get_jumpstart_content_bucket() -> Optional[str]:
+        """Returns JumpStart content bucket."""
+        return JumpStartModelsAccessor._content_bucket
+
     @staticmethod
     def _validate_and_mutate_region_cache_kwargs(
         cache_kwargs: Optional[Dict[str, Any]] = None, region: Optional[str] = None

src/sagemaker/jumpstart/utils.py

@@ -72,21 +72,37 @@ def get_jumpstart_content_bucket(
         RuntimeError: If JumpStart is not launched in ``region``.
     """
 
+    old_content_bucket: Optional[
+        str
+    ] = accessors.JumpStartModelsAccessor.get_jumpstart_content_bucket()
+
+    info_logs: List[str] = []
+
+    bucket_to_return: Optional[str] = None
     if (
         constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE in os.environ
         and len(os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]) > 0
     ):
-        bucket_override = os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]
-        constants.JUMPSTART_LOGGER.info("Using JumpStart bucket override: '%s'", bucket_override)
-        return bucket_override
-    try:
-        return constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[region].content_bucket
-    except KeyError:
-        formatted_launched_regions_str = get_jumpstart_launched_regions_message()
-        raise ValueError(
-            f"Unable to get content bucket for JumpStart in {region} region. "
-            f"{formatted_launched_regions_str}"
-        )
+        bucket_to_return = os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]
+        info_logs.append(f"Using JumpStart bucket override: '{bucket_to_return}'")
+    else:
+        try:
+            bucket_to_return = constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[
+                region
+            ].content_bucket
+        except KeyError:
+            formatted_launched_regions_str = get_jumpstart_launched_regions_message()
+            raise ValueError(
+                f"Unable to get content bucket for JumpStart in {region} region. "
+                f"{formatted_launched_regions_str}"
+            )
+
+    accessors.JumpStartModelsAccessor.set_jumpstart_content_bucket(bucket_to_return)
+
+    if bucket_to_return != old_content_bucket:
+        for info_log in info_logs:
+            constants.JUMPSTART_LOGGER.info(info_log)
+    return bucket_to_return
 
 
 def get_formatted_manifest(

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -57,10 +57,7 @@ def test_get_jumpstart_content_bucket_override():
         with patch("logging.Logger.info") as mocked_info_log:
             random_region = "random_region"
             assert "some-val" == utils.get_jumpstart_content_bucket(random_region)
-            mocked_info_log.assert_called_once_with(
-                "Using JumpStart bucket override: '%s'",
-                "some-val",
-            )
+            mocked_info_log.assert_called_once_with("Using JumpStart bucket override: 'some-val'")
 
 
 def test_get_jumpstart_launched_regions_message():