Skip to content

Commit 9be072b

Browse files
author
Mohamed Zeidan
committed
fix: lazily initialize DEFAULT_JUMPSTART_SAGEMAKER_SESSION (#4468)
v2 maintenance backport. 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.
1 parent 7412b08 commit 9be072b

2 files changed

Lines changed: 155 additions & 11 deletions

File tree

‎src/sagemaker/jumpstart/constants.py‎

Lines changed: 69 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
# ANY KIND, either express or implied. See the License for the specific
1212
# language governing permissions and limitations under the License.
1313
"""This module stores constants related to SageMaker JumpStart."""
14+
1415
from __future__ import absolute_import
16+
import copy
1517
import logging
1618
import os
1719
from typing import Dict, Set, Type
@@ -35,7 +37,6 @@
3537
)
3638
from sagemaker.session import Session
3739

38-
3940
JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart")
4041

4142
# disable logging if env var is set
@@ -176,16 +177,73 @@ def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]:
176177

177178
MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html"
178179

179-
try:
180-
DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(
181-
boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)
182-
)
183-
except Exception as e: # pylint: disable=W0703
184-
DEFAULT_JUMPSTART_SAGEMAKER_SESSION = None
185-
JUMPSTART_LOGGER.warning(
186-
"Unable to create default JumpStart SageMaker Session due to the following error: %s.",
187-
str(e),
188-
)
180+
181+
class _LazyJumpStartSagemakerSession: # pylint: disable=too-few-public-methods
182+
"""Lazily-initialized proxy for the default JumpStart SageMaker Session.
183+
184+
Constructing a real ``Session`` eagerly creates several boto3 clients and
185+
resolves AWS credentials/region. Doing that work at import time added
186+
several seconds to ``import sagemaker`` even when the default session was
187+
never used (see aws/sagemaker-python-sdk#4468).
188+
189+
This proxy defers the work until the session is first used. It:
190+
191+
* is truthy without initializing, so ``session or DEFAULT_...`` fallbacks
192+
and ``if session:`` checks stay cheap and do not build boto3 clients;
193+
* forwards attribute reads and writes to the underlying real session;
194+
* forwards ``copy``/``deepcopy`` so callers that copy-and-mutate the default
195+
session keep working;
196+
* on the failure path (``Session(...)`` construction raises), logs a warning
197+
and resolves to ``None`` so that *attribute access* then behaves exactly as
198+
it would on ``None`` (an ``AttributeError``). Note the degradation is only
199+
observed through attribute access: because the proxy is always truthy,
200+
``bool()``/``if session:``/``or``/``is None`` checks do not see it. In
201+
practice this path is effectively unreachable -- the region is a hardcoded
202+
constant and boto3 defers credential resolution to first API call, so
203+
constructing the clients does not raise at import time.
204+
"""
205+
206+
__slots__ = ()
207+
208+
_resolved = False
209+
_session = None
210+
211+
@classmethod
212+
def _resolve(cls):
213+
"""Build the real Session once, caching the result (or ``None``)."""
214+
if not cls._resolved:
215+
cls._resolved = True
216+
try:
217+
cls._session = Session(boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME))
218+
except Exception as e: # pylint: disable=W0703
219+
cls._session = None
220+
JUMPSTART_LOGGER.warning(
221+
"Unable to create default JumpStart SageMaker Session due to the "
222+
"following error: %s.",
223+
str(e),
224+
)
225+
return cls._session
226+
227+
def __bool__(self):
228+
# Truthy without initializing, so ``x or DEFAULT_...`` stays lazy.
229+
return True
230+
231+
def __getattr__(self, name):
232+
# Delegates to the real session; on the failure path this raises the
233+
# same AttributeError that accessing an attribute on ``None`` would.
234+
return getattr(type(self)._resolve(), name)
235+
236+
def __setattr__(self, name, value):
237+
setattr(type(self)._resolve(), name, value)
238+
239+
def __copy__(self):
240+
return copy.copy(type(self)._resolve())
241+
242+
def __deepcopy__(self, memo):
243+
return copy.deepcopy(type(self)._resolve(), memo)
244+
245+
246+
DEFAULT_JUMPSTART_SAGEMAKER_SESSION = _LazyJumpStartSagemakerSession()
189247

190248
EXTRA_MODEL_ID_TAGS = ["sm-jumpstart-id", "sagemaker-studio:jumpstart-model-id"]
191249
EXTRA_MODEL_VERSION_TAGS = [
Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Tests for lazy initialization of DEFAULT_JUMPSTART_SAGEMAKER_SESSION (GH #4468)."""
14+
15+
from __future__ import absolute_import
16+
17+
import copy
18+
from types import SimpleNamespace
19+
20+
import pytest
21+
from mock import MagicMock, patch
22+
23+
from sagemaker.jumpstart import constants
24+
from sagemaker.jumpstart.constants import _LazyJumpStartSagemakerSession
25+
26+
27+
@pytest.fixture(autouse=True)
28+
def reset_lazy_session_cache():
29+
"""Ensure each test starts and ends with an unresolved proxy cache."""
30+
_LazyJumpStartSagemakerSession._resolved = False
31+
_LazyJumpStartSagemakerSession._session = None
32+
yield
33+
_LazyJumpStartSagemakerSession._resolved = False
34+
_LazyJumpStartSagemakerSession._session = None
35+
36+
37+
def test_default_session_is_a_lazy_proxy():
38+
assert isinstance(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION, _LazyJumpStartSagemakerSession)
39+
40+
41+
def test_truthiness_does_not_build_a_session():
42+
"""``session or DEFAULT_...`` / ``if session:`` must stay lazy (no boto clients)."""
43+
with patch.object(constants, "Session") as session_cls:
44+
assert bool(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION) is True
45+
assert _LazyJumpStartSagemakerSession._resolved is False
46+
session_cls.assert_not_called()
47+
48+
49+
def test_first_attribute_access_builds_session_once():
50+
fake = MagicMock()
51+
fake.boto_region_name = "us-west-2"
52+
with patch.object(constants, "Session", return_value=fake) as session_cls:
53+
assert constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name == "us-west-2"
54+
assert _LazyJumpStartSagemakerSession._resolved is True
55+
_ = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name
56+
session_cls.assert_called_once()
57+
58+
59+
def test_setattr_is_forwarded_to_real_session():
60+
fake = MagicMock()
61+
with patch.object(constants, "Session", return_value=fake):
62+
constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.sagemaker_client = "client"
63+
assert fake.sagemaker_client == "client"
64+
65+
66+
def test_copy_returns_the_real_session():
67+
"""utils.get_default_jumpstart_session_with_user_agent_suffix copies then mutates,
68+
so copy.copy(proxy) must yield a real (copyable, mutable) session, not the proxy."""
69+
fake = SimpleNamespace(boto_session="orig", sagemaker_client="orig")
70+
with patch.object(constants, "Session", return_value=fake):
71+
result = copy.copy(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION)
72+
assert not isinstance(result, _LazyJumpStartSagemakerSession)
73+
assert isinstance(result, SimpleNamespace)
74+
assert result is not fake
75+
result.boto_session = "new"
76+
assert fake.boto_session == "orig"
77+
78+
79+
def test_failed_build_degrades_to_none_contract():
80+
"""If Session construction raises, resolution yields None and logs a warning;
81+
attribute access then behaves exactly as it would on ``None``."""
82+
with patch.object(constants, "Session", side_effect=RuntimeError("boom")):
83+
assert bool(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION) is True
84+
assert _LazyJumpStartSagemakerSession._resolve() is None
85+
with pytest.raises(AttributeError):
86+
_ = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name

0 commit comments

Comments
 (0)