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/unit/sagemaker/jumpstart/test_constants.py

@@ -0,0 +1,141 @@
+# 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():
+    """Start each test with an unresolved proxy cache, then restore what we found.
+
+    The cache lives on the class, i.e. it is process-wide. Blanking it on teardown
+    would discard a session that earlier tests in the same worker had already
+    resolved and mutated, so save and restore the original values instead.
+    """
+    saved_resolved = _LazyJumpStartSagemakerSession._resolved
+    saved_session = _LazyJumpStartSagemakerSession._session
+    _LazyJumpStartSagemakerSession._resolved = False
+    _LazyJumpStartSagemakerSession._session = None
+    try:
+        yield
+    finally:
+        _LazyJumpStartSagemakerSession._resolved = saved_resolved
+        _LazyJumpStartSagemakerSession._session = saved_session
+
+
+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"