From a28524a8635e04f758cae3754f00fe86c1ad3732 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Fri, 25 Sep 2026 12:27:58 -0700 Subject: [PATCH 1/2] fix: lazily initialize DEFAULT_JUMPSTART_SAGEMAKER_SESSION (#4468) Constructing the default JumpStart Session eagerly built ~6 boto3 clients and resolved credentials at import time, adding several seconds to import sagemaker even when the default session was never used. Replace it with a lazy proxy that defers Session construction until first use, stays truthy without initializing, forwards attribute reads/writes and copy/deepcopy, and preserves the historical fail-to-None contract. --- .../src/sagemaker/core/jumpstart/constants.py | 78 ++++++++++++++-- .../tests/unit/jumpstart/test_constants.py | 92 +++++++++++++++++++ 2 files changed, 160 insertions(+), 10 deletions(-) create mode 100644 sagemaker-core/tests/unit/jumpstart/test_constants.py diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/constants.py b/sagemaker-core/src/sagemaker/core/jumpstart/constants.py index 1701095e97..9dd592cebc 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/constants.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/constants.py @@ -14,6 +14,7 @@ from __future__ import absolute_import from __future__ import absolute_import +import copy import logging import os from typing import Dict, Set, Type @@ -180,16 +181,73 @@ 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 __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/sagemaker-core/tests/unit/jumpstart/test_constants.py b/sagemaker-core/tests/unit/jumpstart/test_constants.py new file mode 100644 index 0000000000..02bde92ee3 --- /dev/null +++ b/sagemaker-core/tests/unit/jumpstart/test_constants.py @@ -0,0 +1,92 @@ +# 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 unittest.mock import MagicMock, patch + +from sagemaker.core.jumpstart import constants +from sagemaker.core.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: + # First access materializes the real Session. + assert constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name == "us-west-2" + assert _LazyJumpStartSagemakerSession._resolved is True + # Second access reuses the cached Session (not rebuilt). + _ = 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) + # A shallow copy of the resolved session -- distinct object, same contents, + # and crucially NOT the lazy proxy. + assert not isinstance(result, _LazyJumpStartSagemakerSession) + assert isinstance(result, SimpleNamespace) + assert result is not fake + # Mutating the copy (as the real caller does) must not raise. + result.boto_session = "new" + assert fake.boto_session == "orig" + + +def test_failed_build_degrades_to_none_contract(caplog): + """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")): + # Truthiness is still cheap and does not raise. + 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 From e270d1eaf0ee5bf22163b455ecac853a3c3e87f3 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Sun, 27 Sep 2026 16:52:39 -0700 Subject: [PATCH 2/2] fix: forward __delattr__ on the lazy JumpStart session proxy _LazyJumpStartSagemakerSession declares __slots__ = () and forwards __getattr__/__setattr__ but not __delattr__. unittest.mock.patch reads the original via target.__dict__[name], which the proxy forwards to the real session's instance __dict__; for a class-level attribute such as Session.read_s3_file that raises KeyError, so mock records is_local=False and restores the attribute by calling delattr on teardown. Without __delattr__ that teardown raises AttributeError: '_LazyJumpStartSagemakerSession' object has no attribute ... and, worse, leaves the mock installed on the process-wide session, so unrelated tests in the same worker then see the mock instead of the real attribute. Forwarding __delattr__ removes the shadowing instance attribute created by patch's setattr, which makes the class-level attribute visible again. __slots__ is kept: dropping it does not fix this, because target.__dict__ would then resolve to the proxy's own empty dict and mock would still take the delattr branch. Adds two regression tests covering the class-level path (the defect) and the instance-level path (is_local=True, restored via setattr). --- .../src/sagemaker/core/jumpstart/constants.py | 9 ++++ .../tests/unit/jumpstart/test_constants.py | 46 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/constants.py b/sagemaker-core/src/sagemaker/core/jumpstart/constants.py index 9dd592cebc..b090270ecc 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/constants.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/constants.py @@ -240,6 +240,15 @@ def __getattr__(self, 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()) diff --git a/sagemaker-core/tests/unit/jumpstart/test_constants.py b/sagemaker-core/tests/unit/jumpstart/test_constants.py index 02bde92ee3..2d6b1f52b5 100644 --- a/sagemaker-core/tests/unit/jumpstart/test_constants.py +++ b/sagemaker-core/tests/unit/jumpstart/test_constants.py @@ -90,3 +90,49 @@ def test_failed_build_degrades_to_none_contract(caplog): 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"