diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index b81f97ce3a..043e8ddd07 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/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 = [ diff --git a/tests/conftest.py b/tests/conftest.py index 7839c97eba..350a537222 100644 --- a/tests/conftest.py +++ b/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") diff --git a/tests/data/spark/code/java/hello-java-spark/com/amazonaws/sagemaker/spark/test/HelloJavaSparkApp.class b/tests/data/spark/code/java/hello-java-spark/com/amazonaws/sagemaker/spark/test/HelloJavaSparkApp.class new file mode 100644 index 0000000000..66c6e9ab84 Binary files /dev/null and b/tests/data/spark/code/java/hello-java-spark/com/amazonaws/sagemaker/spark/test/HelloJavaSparkApp.class differ diff --git a/tests/data/spark/code/java/hello-java-spark/hello-spark-java.jar b/tests/data/spark/code/java/hello-java-spark/hello-spark-java.jar new file mode 100644 index 0000000000..827cd79f77 Binary files /dev/null and b/tests/data/spark/code/java/hello-java-spark/hello-spark-java.jar differ diff --git a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py index 8a19baaf40..7341d3ad6c 100644 --- a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py +++ b/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"] diff --git a/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py b/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py index a6e33f1bdf..d512915343 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py +++ b/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( diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 76334330f5..3956c2240d 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/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() diff --git a/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py b/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py index 7cb0f34fbf..429a94f8db 100644 --- a/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py +++ b/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...") diff --git a/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py b/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py index 329b069fee..5b4c89c5b9 100644 --- a/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py +++ b/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...") diff --git a/tests/integ/test_training_compiler.py b/tests/integ/test_training_compiler.py index 1251eb0723..f4fe7d293e 100644 --- a/tests/integ/test_training_compiler.py +++ b/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") diff --git a/tests/unit/sagemaker/jumpstart/test_constants.py b/tests/unit/sagemaker/jumpstart/test_constants.py new file mode 100644 index 0000000000..41e87080a3 --- /dev/null +++ b/tests/unit/sagemaker/jumpstart/test_constants.py @@ -0,0 +1,132 @@ +# 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. +"""Tests for lazy initialization of DEFAULT_JUMPSTART_SAGEMAKER_SESSION (GH #4468).""" + +from __future__ import absolute_import + +import copy +from types import SimpleNamespace + +import pytest +from mock import MagicMock, patch + +from sagemaker.jumpstart import constants +from sagemaker.jumpstart.constants import _LazyJumpStartSagemakerSession + + +@pytest.fixture(autouse=True) +def reset_lazy_session_cache(): + """Ensure each test starts and ends with an unresolved proxy cache.""" + _LazyJumpStartSagemakerSession._resolved = False + _LazyJumpStartSagemakerSession._session = None + yield + _LazyJumpStartSagemakerSession._resolved = False + _LazyJumpStartSagemakerSession._session = None + + +def test_default_session_is_a_lazy_proxy(): + assert isinstance(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION, _LazyJumpStartSagemakerSession) + + +def test_truthiness_does_not_build_a_session(): + """``session or DEFAULT_...`` / ``if session:`` must stay lazy (no boto clients).""" + with patch.object(constants, "Session") as session_cls: + assert bool(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION) is True + assert _LazyJumpStartSagemakerSession._resolved is False + session_cls.assert_not_called() + + +def test_first_attribute_access_builds_session_once(): + fake = MagicMock() + fake.boto_region_name = "us-west-2" + with patch.object(constants, "Session", return_value=fake) as session_cls: + assert constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name == "us-west-2" + assert _LazyJumpStartSagemakerSession._resolved is True + _ = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name + session_cls.assert_called_once() + + +def test_setattr_is_forwarded_to_real_session(): + fake = MagicMock() + with patch.object(constants, "Session", return_value=fake): + constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.sagemaker_client = "client" + assert fake.sagemaker_client == "client" + + +def test_copy_returns_the_real_session(): + """utils.get_default_jumpstart_session_with_user_agent_suffix copies then mutates, + so copy.copy(proxy) must yield a real (copyable, mutable) session, not the proxy.""" + fake = SimpleNamespace(boto_session="orig", sagemaker_client="orig") + with patch.object(constants, "Session", return_value=fake): + result = copy.copy(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION) + assert not isinstance(result, _LazyJumpStartSagemakerSession) + assert isinstance(result, SimpleNamespace) + assert result is not fake + result.boto_session = "new" + assert fake.boto_session == "orig" + + +def test_failed_build_degrades_to_none_contract(): + """If Session construction raises, resolution yields None and logs a warning; + attribute access then behaves exactly as it would on ``None``.""" + with patch.object(constants, "Session", side_effect=RuntimeError("boom")): + assert bool(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION) is True + assert _LazyJumpStartSagemakerSession._resolve() is None + with pytest.raises(AttributeError): + _ = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name + + +def test_mock_patch_of_class_level_attribute_tears_down_cleanly(): + """``mock.patch`` on a class-level session method must restore the original and + leave nothing behind on the process-wide session. + + mock records ``is_local=False`` for a class-level attribute (it is absent from + the instance ``__dict__``) and restores it by calling ``delattr``, so the proxy + must forward ``__delattr__``. Without that forwarding the teardown raises + ``AttributeError`` and the mock leaks into every later test in the same worker. + """ + + class FakeSession: + """Stands in for ``Session``: ``read_s3_file`` is a class-level attribute.""" + + def read_s3_file(self): + return "real" + + fake = FakeSession() + with patch.object(constants, "Session", return_value=fake): + proxy = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION + assert proxy.read_s3_file() == "real" + + with patch.object(proxy, "read_s3_file", return_value="mocked"): + assert proxy.read_s3_file() == "mocked" + + # Teardown must restore the class method and leave no shadowing instance + # attribute behind on the shared session. + assert proxy.read_s3_file() == "real" + assert "read_s3_file" not in fake.__dict__ + + +def test_mock_patch_of_instance_level_attribute_tears_down_cleanly(): + """The instance-attribute path (``is_local=True``, restored via ``setattr``) + must keep working -- guards against a regression in ``__setattr__`` forwarding.""" + + class FakeSession: + def __init__(self): + self.sagemaker_client = "real-client" + + fake = FakeSession() + with patch.object(constants, "Session", return_value=fake): + proxy = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION + with patch.object(proxy, "sagemaker_client", "mock-client"): + assert proxy.sagemaker_client == "mock-client" + assert proxy.sagemaker_client == "real-client"