fix: lazily initialize DEFAULT_JUMPSTART_SAGEMAKER_SESSION (#4468) [v2]

src/sagemaker/jumpstart/constants.py

@@ -11,7 +11,9 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """This module stores constants related to SageMaker JumpStart."""
+
 from __future__ import absolute_import
+import copy
 import logging
 import os
 from typing import Dict, Set, Type
@@ -35,7 +37,6 @@
 )
 from sagemaker.session import Session
 
-
 JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart")
 
 # disable logging if env var is set
@@ -176,16 +177,82 @@ def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]:
 
 MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html"
 
-try:
-    DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(
-        boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)
-    )
-except Exception as e:  # pylint: disable=W0703
-    DEFAULT_JUMPSTART_SAGEMAKER_SESSION = None
-    JUMPSTART_LOGGER.warning(
-        "Unable to create default JumpStart SageMaker Session due to the following error: %s.",
-        str(e),
-    )
+
+class _LazyJumpStartSagemakerSession:  # pylint: disable=too-few-public-methods
+    """Lazily-initialized proxy for the default JumpStart SageMaker Session.
+
+    Constructing a real ``Session`` eagerly creates several boto3 clients and
+    resolves AWS credentials/region. Doing that work at import time added
+    several seconds to ``import sagemaker`` even when the default session was
+    never used (see aws/sagemaker-python-sdk#4468).
+
+    This proxy defers the work until the session is first used. It:
+
+    * is truthy without initializing, so ``session or DEFAULT_...`` fallbacks
+      and ``if session:`` checks stay cheap and do not build boto3 clients;
+    * forwards attribute reads and writes to the underlying real session;
+    * forwards ``copy``/``deepcopy`` so callers that copy-and-mutate the default
+      session keep working;
+    * on the failure path (``Session(...)`` construction raises), logs a warning
+      and resolves to ``None`` so that *attribute access* then behaves exactly as
+      it would on ``None`` (an ``AttributeError``). Note the degradation is only
+      observed through attribute access: because the proxy is always truthy,
+      ``bool()``/``if session:``/``or``/``is None`` checks do not see it. In
+      practice this path is effectively unreachable -- the region is a hardcoded
+      constant and boto3 defers credential resolution to first API call, so
+      constructing the clients does not raise at import time.
+    """
+
+    __slots__ = ()
+
+    _resolved = False
+    _session = None
+
+    @classmethod
+    def _resolve(cls):
+        """Build the real Session once, caching the result (or ``None``)."""
+        if not cls._resolved:
+            cls._resolved = True
+            try:
+                cls._session = Session(boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME))
+            except Exception as e:  # pylint: disable=W0703
+                cls._session = None
+                JUMPSTART_LOGGER.warning(
+                    "Unable to create default JumpStart SageMaker Session due to the "
+                    "following error: %s.",
+                    str(e),
+                )
+        return cls._session
+
+    def __bool__(self):
+        # Truthy without initializing, so ``x or DEFAULT_...`` stays lazy.
+        return True
+
+    def __getattr__(self, name):
+        # Delegates to the real session; on the failure path this raises the
+        # same AttributeError that accessing an attribute on ``None`` would.
+        return getattr(type(self)._resolve(), name)
+
+    def __setattr__(self, name, value):
+        setattr(type(self)._resolve(), name, value)
+
+    def __delattr__(self, name):
+        # Forwarded so that ``unittest.mock.patch`` tears down cleanly. For a
+        # class-level attribute mock records ``is_local=False``, sets the mock
+        # via ``__setattr__`` (creating an instance attribute on the real
+        # session) and restores by calling ``delattr``. Without this method the
+        # teardown raises ``AttributeError`` and leaves the mock installed on
+        # the process-wide session.
+        delattr(type(self)._resolve(), name)
+
+    def __copy__(self):
+        return copy.copy(type(self)._resolve())
+
+    def __deepcopy__(self, memo):
+        return copy.deepcopy(type(self)._resolve(), memo)
+
+
+DEFAULT_JUMPSTART_SAGEMAKER_SESSION = _LazyJumpStartSagemakerSession()
 
 EXTRA_MODEL_ID_TAGS = ["sm-jumpstart-id", "sagemaker-studio:jumpstart-model-id"]
 EXTRA_MODEL_VERSION_TAGS = [

tests/conftest.py

@@ -641,42 +641,17 @@ def cpu_instance_type(sagemaker_session, request):
 
 @pytest.fixture(scope="session")
 def gpu_instance_type(sagemaker_session, request):
-    region = sagemaker_session.boto_session.region_name
-    if region in NO_P3_REGIONS:
-        return "ml.p2.xlarge"
-    else:
-        return "ml.p3.2xlarge"
+    return "ml.g4dn.xlarge"
 
 
 @pytest.fixture()
 def gpu_pytorch_instance_type(sagemaker_session, request):
-    fw_version = None
-    for pytorch_version_fixture in [
-        "pytorch_inference_version",
-        "huggingface_training_pytorch_latest_version",
-        "huggingface_inference_pytorch_latest_version",
-    ]:
-        if pytorch_version_fixture in request.fixturenames:
-            fw_version = request.getfixturevalue(pytorch_version_fixture)
-    if fw_version is None:
-        fw_version = request.param
-    region = sagemaker_session.boto_session.region_name
-    if region in NO_P3_REGIONS:
-        if Version(fw_version) >= Version("1.13"):
-            return PYTORCH_RENEWED_GPU
-        else:
-            return "ml.p2.xlarge"
-    else:
-        return "ml.p3.2xlarge"
+    return "ml.g4dn.xlarge"
 
 
 @pytest.fixture(scope="session")
 def gpu_instance_type_list(sagemaker_session, request):
-    region = sagemaker_session.boto_session.region_name
-    if region in NO_P3_REGIONS:
-        return ["ml.p2.xlarge"]
-    else:
-        return ["ml.p3.2xlarge", "ml.p2.xlarge"]
+    return ["ml.g4dn.xlarge"]
 
 
 @pytest.fixture(scope="session")
@@ -717,16 +692,7 @@ def pytest_generate_tests(metafunc):
         cpu_instance_type = "ml.m5.xlarge" if region in NO_M4_REGIONS else "ml.m4.xlarge"
 
         params = [cpu_instance_type]
-        if not (
-            region in tests.integ.HOSTING_NO_P3_REGIONS
-            or region in tests.integ.TRAINING_NO_P3_REGIONS
-        ):
-            params.append("ml.p3.2xlarge")
-        elif not (
-            region in tests.integ.HOSTING_NO_P2_REGIONS
-            or region in tests.integ.TRAINING_NO_P2_REGIONS
-        ):
-            params.append("ml.p2.xlarge")
+        params.append("ml.g4dn.xlarge")
 
         metafunc.parametrize("instance_type", params, scope="session")
 

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -95,9 +95,10 @@ def test_prepacked_jumpstart_model(setup):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.p3.2xlarge instance
+    # uses ml.g4dn.xlarge instance
     predictor = model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g4dn.xlarge",
     )
 
     response = predictor.predict("hello world!")
@@ -120,7 +121,7 @@ def test_model_package_arn_jumpstart_model(setup):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.g5.2xlarge instance
+    # uses ml.g4dn.2xlarge instance
     predictor = model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
     )
@@ -175,7 +176,7 @@ def test_jumpstart_gated_model(setup):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.g5.2xlarge instance
+    # uses ml.g4dn.2xlarge instance
     predictor = model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         accept_eula=True,
@@ -206,6 +207,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup):
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         accept_eula=True,
         endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED,
+        instance_type="ml.g5.2xlarge",
     )
 
     predictor = retrieve_default(
@@ -262,7 +264,7 @@ def test_jumpstart_model_register(setup):
 
     # uses  instance
     predictor = model_package.deploy(
-        instance_type="ml.p3.2xlarge",
+        instance_type="ml.g4dn.xlarge",
         initial_instance_count=1,
     )
 
@@ -379,7 +381,7 @@ def test_jumpstart_model_with_deployment_configs(setup):
 
     model.set_deployment_config(
         configs[0]["ConfigName"],
-        "ml.g5.2xlarge",
+        "ml.g4dn.2xlarge",
     )
     assert model.config_name == configs[0]["ConfigName"]
 

tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py

@@ -69,6 +69,7 @@ def test_jumpstart_hub_estimator(setup, add_model_references):
         model_id=model_id,
         hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g4dn.xlarge",
     )
 
     estimator.fit(
@@ -85,9 +86,10 @@ def test_jumpstart_hub_estimator(setup, add_model_references):
         model_version=model_version,
     )
 
-    # uses ml.p3.2xlarge instance
+    # uses ml.g4dn.xlarge instance
     predictor = estimator.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g4dn.xlarge",
     )
 
     response = predictor.predict(["hello", "world"])
@@ -107,6 +109,7 @@ def test_jumpstart_hub_estimator_with_session(setup, add_model_references):
         sagemaker_session=sagemaker_session,
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
+        instance_type="ml.g4dn.xlarge",
     )
 
     estimator.fit(
@@ -124,11 +127,12 @@ def test_jumpstart_hub_estimator_with_session(setup, add_model_references):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.p3.2xlarge instance
+    # uses ml.g4dn.xlarge instance
     predictor = estimator.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
+        instance_type="ml.g4dn.xlarge",
     )
 
     response = predictor.predict(["hello", "world"])
@@ -144,6 +148,7 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references):
         model_id=model_id,
         hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g5.2xlarge",
     )
 
     estimator.fit(
@@ -158,6 +163,7 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references):
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
+        instance_type="ml.g5.2xlarge",
     )
 
     payload = {
@@ -178,6 +184,7 @@ def test_jumpstart_hub_gated_estimator_without_eula(setup, add_model_references)
         model_id=model_id,
         hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g5.2xlarge",
     )
     with pytest.raises(Exception):
         estimator.fit(

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

@@ -116,6 +116,7 @@ def test_jumpstart_hub_gated_model(setup, add_model_references):
     predictor = model.deploy(
         accept_eula=True,
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g5.2xlarge",
     )
 
     payload = model.retrieve_example_payload()

tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py

@@ -34,8 +34,8 @@
 def test_jumpstart_transfer_learning_retrieve_functions(setup):
 
     model_id, model_version = "huggingface-spc-bert-base-cased", "2.0.3"
-    training_instance_type = "ml.p3.2xlarge"
-    inference_instance_type = "ml.p2.xlarge"
+    training_instance_type = "ml.g4dn.xlarge"
+    inference_instance_type = "ml.g4dn.xlarge"
 
     # training
     print("Starting training...")

tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py

@@ -53,12 +53,8 @@ def test_jumpstart_transfer_learning_estimator_class(setup):
 
     model_id, model_version = "huggingface-spc-bert-base-cased", "1.2.3"
 
-    inference_instance_type = instance_types.retrieve_default(
-        model_id=model_id, model_version=model_version, scope="inference"
-    )
-    training_instance_type = instance_types.retrieve_default(
-        model_id=model_id, model_version=model_version, scope="training"
-    )
+    inference_instance_type = "ml.g4dn.xlarge"
+    training_instance_type = "ml.g4dn.xlarge"
     instance_count = 1
 
     print("Starting training...")

tests/integ/test_training_compiler.py

@@ -30,7 +30,7 @@
 
 @pytest.fixture(scope="module")
 def gpu_instance_type(request):
-    return "ml.p3.2xlarge"
+    return "ml.g4dn.xlarge"
 
 
 @pytest.fixture(scope="module")
@@ -86,8 +86,8 @@ def skip_if_incompatible(gpu_instance_type, request):
 @pytest.mark.parametrize(
     "gpu_instance_type,instance_count",
     [
-        pytest.param("ml.p3.2xlarge", 1, marks=pytest.mark.release),
-        pytest.param("ml.p3.16xlarge", 2),
+        pytest.param("ml.g4dn.xlarge", 1, marks=pytest.mark.release),
+        pytest.param("ml.g4dn.12xlarge", 2),
     ],
 )
 @pytest.mark.skipif(
@@ -138,8 +138,8 @@ def test_huggingface_pytorch(
 @pytest.mark.parametrize(
     "gpu_instance_type,instance_count",
     [
-        pytest.param("ml.p3.2xlarge", 1, marks=pytest.mark.release),
-        pytest.param("ml.p3.16xlarge", 2),
+        pytest.param("ml.g4dn.xlarge", 1, marks=pytest.mark.release),
+        pytest.param("ml.g4dn.12xlarge", 2),
     ],
 )
 @pytest.mark.skip("Temporarily skip to unblock")