fix: lazily initialize DEFAULT_JUMPSTART_SAGEMAKER_SESSION (#4468) [v2] - #6331
mohamedzeidan2021 wants to merge 2 commits into
Conversation
04b9546 to
d16e4ed
Compare
jam-jee
left a comment
jam-jee
left a comment
There was a problem hiding this comment.
Summary: The fix itself is fine — src/sagemaker/jumpstart/constants.py is the same proxy as #6330 and the v2 copy(...) call site at utils.py:1304 is handled the same way. The branch, however, is not mergeable as-is.
Blocking
-
GitHub reports
CONFLICTING/DIRTY. The branch carries two unrelated commits dated 2026-05-01:18f22415 "v2 test updates"(rewritestests/conftest.pyGPU fixtures toml.g4dn.xlarge, edits 6 integ tests, and adds two binarieshello-spark-java.jar/HelloJavaSparkApp.class) and7412b08e "instance types update". They leak 9 unrelated files into this PR (+240/−69 vs the ~+70/−10 the fix needs) and overlap88dbcc2e fix(tests): Replace deprecated ml.t2.medium (#6284)already onmaster-v2, which is the conflict.master-v2...d16e4edis ahead 4 / behind 27.Please
git rebase --onto origin/master-v2 7412b08eso only9be072b4+d16e4ederemain. The resulting diff should be exactlysrc/sagemaker/jumpstart/constants.py+tests/unit/sagemaker/jumpstart/test_constants.py. -
The
docs/readthedocs.org:sagemakerfailure (No module named 'pkg_resources') is because this base predates570b04c4 fix: RTD build failure (#6023); the rebase fixes it (v2 PR #6345 built today with the same setuptools and passed).
Should-have
- Carry over
test_mock_patch_of_instance_level_attribute_tears_down_cleanlyfrom #6330 — v2 has many tests thatpatch.objectthe default session, so that teardown path matters more here than on v3.
Nits (same as #6330)
_resolveis not thread-safe; two threads racing on first use both build aSessionand one is discarded (the old module-level constant was built under the import lock). Athreading.Lockis a few lines._resolved = Trueis set before thetry; aKeyboardInterruptmid-construction leaves the proxy permanently resolved toNone. Set it inside both branches.- Add a
__repr__; dunders bypass__getattr__, so Sphinx will render<_LazyJumpStartSagemakerSession object at 0x…>in everysagemaker_session=DEFAULT_...signature.
CI attribution: codestyle-doc-tests (flake8 E501/F401 and black violations in files this PR does not touch; identical on unrelated v2 PR #6337) is pre-existing on master-v2. Unit tests pass on py39-py312.
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.
_LazyJumpStartSagemakerSession declares __slots__ = () and forwards __getattr__/__setattr__ but not __delattr__. 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 leaves the mock installed on the process-wide session, so unrelated tests in the same xdist worker then see the mock instead of the real attribute. That is what made test_notebook_utils, test_model, test_sagemaker_config and test_js_builder fail on py39-py312. 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).
d16e4ed to
bf0f5f5
Compare
This branch was successfully deployed
2 participants
Issue
Fixes #4468 on the v2 maintenance branch. Companion to #6330 (v3, against
master).At import time
DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(boto3.Session(...))eagerly builds ~6 boto3 clients and resolves credentials/region, adding several seconds toimport sagemakereven when the default session is never used.Fix
Same lazy-proxy approach as #6330, ported to v2:
_LazyJumpStartSagemakerSessiondefersSessionconstruction until first use — truthy without initializing, forwards attribute reads/writes andcopy/deepcopy, preserves the fail-to-Nonebehavior on attribute access. The public name and its use as a default argument are unchanged.Validation
boto3.Session.client/.resource: import now builds 0 clients (was 6); first attribute access builds them and yields a workingSession.tests/unit/sagemaker/jumpstart/test_constants.py(6 tests); fails without the fix.black/flake8clean.Backwards compatibility
No public API change; security/bug-fix-only policy respected (behavior for existing inputs is unchanged in the normal path).