Skip to content

Commit 4f3cc56

Browse files
mohamedzeidan2021Mohamed Zeidan
andauthored
fix: lazily initialize DEFAULT_JUMPSTART_SAGEMAKER_SESSION (#4468) (#6330)
* 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. * 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). --------- Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>
1 parent 30c507a commit 4f3cc56

2 files changed

Lines changed: 215 additions & 10 deletions

File tree

‎sagemaker-core/src/sagemaker/core/jumpstart/constants.py‎

Lines changed: 77 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from __future__ import absolute_import
1616
from __future__ import absolute_import
17+
import copy
1718
import logging
1819
import os
1920
from typing import Dict, Set, Type
@@ -180,16 +181,82 @@ def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]:
180181

181182
MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html"
182183

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

194261
EXTRA_MODEL_ID_TAGS = ["sm-jumpstart-id", "sagemaker-studio:jumpstart-model-id"]
195262
EXTRA_MODEL_VERSION_TAGS = [
Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
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 unittest.mock import MagicMock, patch
22+
23+
from sagemaker.core.jumpstart import constants
24+
from sagemaker.core.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+
# First access materializes the real Session.
54+
assert constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name == "us-west-2"
55+
assert _LazyJumpStartSagemakerSession._resolved is True
56+
# Second access reuses the cached Session (not rebuilt).
57+
_ = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name
58+
session_cls.assert_called_once()
59+
60+
61+
def test_setattr_is_forwarded_to_real_session():
62+
fake = MagicMock()
63+
with patch.object(constants, "Session", return_value=fake):
64+
constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.sagemaker_client = "client"
65+
assert fake.sagemaker_client == "client"
66+
67+
68+
def test_copy_returns_the_real_session():
69+
"""utils.get_default_jumpstart_session_with_user_agent_suffix copies then mutates,
70+
so copy.copy(proxy) must yield a real (copyable, mutable) session, not the proxy."""
71+
fake = SimpleNamespace(boto_session="orig", sagemaker_client="orig")
72+
with patch.object(constants, "Session", return_value=fake):
73+
result = copy.copy(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION)
74+
# A shallow copy of the resolved session -- distinct object, same contents,
75+
# and crucially NOT the lazy proxy.
76+
assert not isinstance(result, _LazyJumpStartSagemakerSession)
77+
assert isinstance(result, SimpleNamespace)
78+
assert result is not fake
79+
# Mutating the copy (as the real caller does) must not raise.
80+
result.boto_session = "new"
81+
assert fake.boto_session == "orig"
82+
83+
84+
def test_failed_build_degrades_to_none_contract(caplog):
85+
"""If Session construction raises, resolution yields None and logs a warning;
86+
attribute access then behaves exactly as it would on ``None``."""
87+
with patch.object(constants, "Session", side_effect=RuntimeError("boom")):
88+
# Truthiness is still cheap and does not raise.
89+
assert bool(constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION) is True
90+
assert _LazyJumpStartSagemakerSession._resolve() is None
91+
with pytest.raises(AttributeError):
92+
_ = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.boto_region_name
93+
94+
95+
def test_mock_patch_of_class_level_attribute_tears_down_cleanly():
96+
"""``mock.patch`` on a class-level session method must restore the original and
97+
leave nothing behind on the process-wide session.
98+
99+
mock records ``is_local=False`` for a class-level attribute (it is absent from
100+
the instance ``__dict__``) and restores it by calling ``delattr``, so the proxy
101+
must forward ``__delattr__``. Without that forwarding the teardown raises
102+
``AttributeError`` and the mock leaks into every later test in the same worker.
103+
"""
104+
105+
class FakeSession:
106+
"""Stands in for ``Session``: ``read_s3_file`` is a class-level attribute."""
107+
108+
def read_s3_file(self):
109+
return "real"
110+
111+
fake = FakeSession()
112+
with patch.object(constants, "Session", return_value=fake):
113+
proxy = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION
114+
assert proxy.read_s3_file() == "real"
115+
116+
with patch.object(proxy, "read_s3_file", return_value="mocked"):
117+
assert proxy.read_s3_file() == "mocked"
118+
119+
# Teardown must restore the class method and leave no shadowing instance
120+
# attribute behind on the shared session.
121+
assert proxy.read_s3_file() == "real"
122+
assert "read_s3_file" not in fake.__dict__
123+
124+
125+
def test_mock_patch_of_instance_level_attribute_tears_down_cleanly():
126+
"""The instance-attribute path (``is_local=True``, restored via ``setattr``)
127+
must keep working -- guards against a regression in ``__setattr__`` forwarding."""
128+
129+
class FakeSession:
130+
def __init__(self):
131+
self.sagemaker_client = "real-client"
132+
133+
fake = FakeSession()
134+
with patch.object(constants, "Session", return_value=fake):
135+
proxy = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION
136+
with patch.object(proxy, "sagemaker_client", "mock-client"):
137+
assert proxy.sagemaker_client == "mock-client"
138+
assert proxy.sagemaker_client == "real-client"

0 commit comments

Comments
 (0)