fix: lazily initialize DEFAULT_JUMPSTART_SAGEMAKER_SESSION (#4468) [v2] - #6331
Open
mohamedzeidan2021 wants to merge 4 commits into
Open
mohamedzeidan2021 wants to merge 4 commits into
mohamedzeidan2021 wants to merge 4 commits into
Conversation
added 3 commits
April 30, 2026 18:58
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).
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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).