feat: override jumpstart content bucket (#2901)

Co-authored-by: Navin Soni <navinsoni89@gmail.com>
Co-authored-by: Mufaddal Rohawala <89424143+mufaddal-rohawala@users.noreply.github.com>
Co-authored-by: qidewenwhen <32910701+qidewenwhen@users.noreply.github.com>
Co-authored-by: HappyAmazonian <91216626+HappyAmazonian@users.noreply.github.com>

Author: 5 people

doc/overview.rst

@@ -746,6 +746,7 @@ see `Model <https://sagemaker.readthedocs.io/en/stable/api/inference/model.html
 .. code:: python
 
    from sagemaker.model import Model
+   from sagemaker.predictor import Predictor
    from sagemaker.session import Session
 
    # Create the SageMaker model instance
@@ -755,6 +756,7 @@ see `Model <https://sagemaker.readthedocs.io/en/stable/api/inference/model.html
        source_dir=script_uri,
        entry_point="inference.py",
        role=Session().get_caller_identity_arn(),
+       predictor_cls=Predictor,
    )
 
 Save the output from deploying the model to a variable named
@@ -766,12 +768,9 @@ Deployment may take about 5 minutes.
 
 .. code:: python
 
-   from sagemaker.predictor import Predictor
-
    predictor = model.deploy(
        initial_instance_count=instance_count,
        instance_type=instance_type,
-       predictor_cls=Predictor
    )
 
 Because ``catboost`` and ``lightgbm`` rely on the PyTorch Deep Learning Containers

src/sagemaker/jumpstart/constants.py

@@ -122,3 +122,5 @@
 TRAINING_ENTRY_POINT_SCRIPT_NAME = "transfer_learning.py"
 
 SUPPORTED_JUMPSTART_SCOPES = set(scope.value for scope in JumpStartScriptScope)
+
+ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE = "AWS_JUMPSTART_CONTENT_BUCKET_OVERRIDE"

src/sagemaker/jumpstart/utils.py

@@ -13,6 +13,7 @@
 """This module contains utilities related to SageMaker JumpStart."""
 from __future__ import absolute_import
 import logging
+import os
 from typing import Dict, List, Optional
 from urllib.parse import urlparse
 from packaging.version import Version
@@ -60,6 +61,14 @@ def get_jumpstart_content_bucket(region: str) -> str:
     Raises:
         RuntimeError: If JumpStart is not launched in ``region``.
     """
+
+    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]
+        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:

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -11,11 +11,13 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+import os
 from mock.mock import Mock, patch
 import pytest
 import random
 from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.constants import (
+    ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE,
     JUMPSTART_BUCKET_NAME_SET,
     JUMPSTART_REGION_NAME_SET,
     JumpStartScriptScope,
@@ -40,6 +42,17 @@ def test_get_jumpstart_content_bucket():
         utils.get_jumpstart_content_bucket(bad_region)
 
 
+def test_get_jumpstart_content_bucket_override():
+    with patch.dict(os.environ, {ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE: "some-val"}):
+        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",
+            )
+
+
 def test_get_jumpstart_launched_regions_message():
 
     with patch("sagemaker.jumpstart.constants.JUMPSTART_REGION_NAME_SET", {}):