From 18f22415148bce747b3f06162e2e1ffd57008617 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Thu, 30 Apr 2026 18:58:16 -0700 Subject: [PATCH 1/4] v2 test updates --- tests/conftest.py | 42 ++---------------- .../spark/test/HelloJavaSparkApp.class | Bin 0 -> 495 bytes .../hello-java-spark/hello-spark-java.jar | Bin 0 -> 4781 bytes .../jumpstart/model/test_jumpstart_model.py | 12 ++--- .../test_jumpstart_private_hub_estimator.py | 8 +++- .../retrieve_uri/test_transfer_learning.py | 4 +- .../test_transfer_learning.py | 8 +--- tests/integ/test_training_compiler.py | 10 ++--- 8 files changed, 26 insertions(+), 58 deletions(-) create mode 100644 tests/data/spark/code/java/hello-java-spark/com/amazonaws/sagemaker/spark/test/HelloJavaSparkApp.class create mode 100644 tests/data/spark/code/java/hello-java-spark/hello-spark-java.jar diff --git a/tests/conftest.py b/tests/conftest.py index 7839c97eba..350a537222 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -641,42 +641,17 @@ def cpu_instance_type(sagemaker_session, request): @pytest.fixture(scope="session") def gpu_instance_type(sagemaker_session, request): - region = sagemaker_session.boto_session.region_name - if region in NO_P3_REGIONS: - return "ml.p2.xlarge" - else: - return "ml.p3.2xlarge" + return "ml.g4dn.xlarge" @pytest.fixture() def gpu_pytorch_instance_type(sagemaker_session, request): - fw_version = None - for pytorch_version_fixture in [ - "pytorch_inference_version", - "huggingface_training_pytorch_latest_version", - "huggingface_inference_pytorch_latest_version", - ]: - if pytorch_version_fixture in request.fixturenames: - fw_version = request.getfixturevalue(pytorch_version_fixture) - if fw_version is None: - fw_version = request.param - region = sagemaker_session.boto_session.region_name - if region in NO_P3_REGIONS: - if Version(fw_version) >= Version("1.13"): - return PYTORCH_RENEWED_GPU - else: - return "ml.p2.xlarge" - else: - return "ml.p3.2xlarge" + return "ml.g4dn.xlarge" @pytest.fixture(scope="session") def gpu_instance_type_list(sagemaker_session, request): - region = sagemaker_session.boto_session.region_name - if region in NO_P3_REGIONS: - return ["ml.p2.xlarge"] - else: - return ["ml.p3.2xlarge", "ml.p2.xlarge"] + return ["ml.g4dn.xlarge"] @pytest.fixture(scope="session") @@ -717,16 +692,7 @@ def pytest_generate_tests(metafunc): cpu_instance_type = "ml.m5.xlarge" if region in NO_M4_REGIONS else "ml.m4.xlarge" params = [cpu_instance_type] - if not ( - region in tests.integ.HOSTING_NO_P3_REGIONS - or region in tests.integ.TRAINING_NO_P3_REGIONS - ): - params.append("ml.p3.2xlarge") - elif not ( - region in tests.integ.HOSTING_NO_P2_REGIONS - or region in tests.integ.TRAINING_NO_P2_REGIONS - ): - params.append("ml.p2.xlarge") + params.append("ml.g4dn.xlarge") metafunc.parametrize("instance_type", params, scope="session") diff --git a/tests/data/spark/code/java/hello-java-spark/com/amazonaws/sagemaker/spark/test/HelloJavaSparkApp.class b/tests/data/spark/code/java/hello-java-spark/com/amazonaws/sagemaker/spark/test/HelloJavaSparkApp.class new file mode 100644 index 0000000000000000000000000000000000000000..66c6e9ab84fe054c85f3f11b178220b7b9618681 GIT binary patch literal 495 zcmZvZ%Sr<=6o&uw;&jH**4EZ*p0DhY7>Nq4HI7 z#E@;Z`z8u74Oli7VgHdkk8vVBQUxoC9D`I+o)ibp)$vb!3edl>l3Q=!%7cLAmxv{?Y@B;LuVLx zTzKLs@Wn&SV=?vPf%)Eh&uEoHmm7FIO z5o@j6c}ND!31-nyo+>S_{c}FhH$f9sdUM?}VEWZ*T%%c?BBPk!eFJ-@gkhaVy$xhx iVgs9$Yw1o}3Yaf2%1f^({>t@#k_Ua0F5IF#gY9okpm9S0 literal 0 HcmV?d00001 diff --git a/tests/data/spark/code/java/hello-java-spark/hello-spark-java.jar b/tests/data/spark/code/java/hello-java-spark/hello-spark-java.jar new file mode 100644 index 0000000000000000000000000000000000000000..827cd79f77158ba2f7aad2350c80c7877d30ab02 GIT binary patch literal 4781 zcmbVQ2{;sL7amK*SvJNi5f+|mFyAIedTzCd{2gWCC9aikhx{hXVZ#6&7NPABy%mc?|%R=SjHAJ zh@8rymO(INj~BoR(f?V}4YP+i>(AGGGKqQq!4B3PX9_ZDYB;lc<6B^HahgRy zlW0kPMi(SY#y9WIs@jGtCL=XWDa1^?T)i5OCNHi(j-?Qwph&3-fu(AVo~3R)lPMvw zQ)v(c52|cB7Nj(UYuTy5m6v=(_eRDi+>J?&yK5S9tB4!g(G)bFSrIgpspKBg!_&kP z2(o{#rYxRwp`=`qlc8vi+atm%l~Oqjs09{MZZAlYJDF4GCO&YEl4M}Qw~AK5V5!wJ zEz6WmHW5+vt_&<&PBLp*}At@$8y~cE&U2KS1Vj2HU%Y_CRfcGqD4e*hG51|OBcgIa!8L~M zf;(V;i=XmI4$RwRvMFS}b$xBJ)9czUQCItWrDSk`l-|F?pWF{vhv$=zvATizl zp5w&t98f!`hl4%T4e>9$EI;ugpq6kus14ltU!*wxMG8TzuOqT=p8T~K&;M6h7dXP@ zyMn)eg(Uws1wR+Oxh)idu+vrtzn7!;9(j93czAdiXP|uY1|QbcBb$mVP+9xj?UPBR zi?_?v3(pN*#=TunnXKoW3}Y0R*D7ceDH2^9uZC`-*B063DCua<+@qVeoz_0ft2xO2 zii4(k9VYZeK}qWx{X?y%{5jI~xFcA;uqFV4DJa70mPqo2Ws4GhWd&<+WiA zM28o};Wpw;=K82p3m3h!piQwLlP=7PNLGKKh*k`|%)$3UBkE1Y)GerTxK!+U4J~Yt zUvk$TNm$IOM_rCofm@E!6oPGvYx=_ODjY&6HIWx*e%qImIO ziF;Z^$WbCLs+)GuNT|nVefv4dL{5b+2W7%lo_2rcwqD+tW=aXNfbcV&vAm{^4M=&7 zOZA%aEffVQ(c(%24{rgn88ph>s>qjhVKUAw|GeE!|HXkJMsMcz-Jv;1t@Ke5Z1 zhQ1t4y`I-;QsU5$&#~B`Bu1c!)gd&2;9m=29Y)~C+$T2B`@}}X|Gi51iH`lzhrYnf zCe6(5_DBi_kQ{du$rswL47i5%z*Wp2NHx8q2k2gIcvoNl@TD$!;1^YIcKr1*V}FHp zud#bww$qaL&Q5K1BX8^kF?D*_M3s$6wqA{!+1<;^1gc$glQcE9uLXBb`E)(olGH~d zk=u|&i7soq;jj=i?W6g=QmPl9$=O#ow^#i1jl6DmrKcF9E13B&et4Ym0Sq*0>z$W? zo0^!LB);5qT~Z&%l^UFBLmCJ6nN|tDOF*Ibcl6psQnFvLgpZ9u{IT}hXE&8QT3WVL zL10<|Wl$>;#M|Ab3YyvvkL8Zvo~-wO>{ntsnt4JE;qKzTnX9ufa`mmi@_?_b0vVZk zY;Z>t9oPc20=xW?(Ohjff_d3~*}(_@(9dpfb8WX5@1tMs!?Zj7wIo|_rfdJY&i#Iw zW?9Vz_)M2RQ{9vcEQ}*^Z*D!)ye;;kmEEo9W%hS{WPD;w@Ri+RG7xYwKjhQ$aFki? zdpAw2c`QGr1f*cTpn!jHEUK7yGL&?Hp(Q3ZR@O}knT}kNY#pBWSQ}_bmhxR3^q5+G zts-i1OARF}&6W4TyU1{A=oLy(e`%q*qPHYqg%S<7y%sJyOC=R z!*C(Pd{zMtWMQJLz-udShamPoCK0Gr93D4zy)$B2TMe~gQ_;iY|n?P zEQfN>eyXiBSn!48#b=sX*lN~FtXG#fvy;aoPE8!4h0?XQni?xUbU) z(MhCvt-r-HlW)x-rAtohJ^|qHd~0%2@Z-8mn&zbRmj_Y1yhbk#`-?&%u|ZJ-SeM1; zlJ5&!%Ib}ch846;-h~cjDG!(Xqe!CTZ8G4ouEN4XwlMz-7Ro+Kn(ZLD#Kh$6Fr#Cf zx=Jv6*c~5J*aKmX5=Iv5XQ%pf5T}ZCo~_J9zJHa?5?vR*o(kk@CEtimV~(3c3p~jQ z6a|RB=b63044kFFd2k}>-_+EMH?H*O8r5-0Z=;GKO?PrJTG$mTUD~PSvPHXYrl>TBeZel$H1>P zn<;JDWBBS!+Kh)?1SSWflO0Hh#s~RDqXm+djpn}%zE%Khc*aglGqrJ6@OY@@C#5d{ zSF(zr`#Z8B_^DWhs?|&4*SnU1p7tj8Fj*fH8VZwk!>D$AN8hcxM9y~ zO7zGtEkOE@AK$GcF+hZERqPw#BodW1RTU>BsLJ>EN8l{2&ido$?|3JqcSfX1*-&fN zS|xp*PsJUehEJXMCvBUL5?mP#6f2L%&)gD3tN$Cbr9n^}g$K%kGhO=tq|@c>}CtO%VRs&GN}RrH%(C%~NZ22yyBIUS!i z@HwBp{?&WQ6Wg$`)bm1XI~<)FG@4= z#+$l0EH#t%mR_Sf-uGKw;=cWkNTfpQ+mH+UJ9_>_r{nJ!5qo_qepwV0M z;Bj9TkBDa~qv;=C#={zJln*wt`?m60FwtM3F$_sR?pUP5DcH_H@p}v+@PpPuZ_#A9F7%eArm=P>Y%+p=(I~`uBNQtpvx22lCuFv3h^_iu2fDAf38St8gtr2U^{k3x@s)&JoY{GlSk&hfwM`W>1a-U;#Fl%vxx5u6-S!T+Z4=WF{8RerqtqcbLPmOdnF0%H94 zAbdCc=W~by=+R?{@i|&{e@#+bLNWf`_ Date: Fri, 1 May 2026 00:30:50 -0700 Subject: [PATCH 2/4] instance types update --- .../estimator/test_jumpstart_private_hub_estimator.py | 3 +++ .../private_hub/model/test_jumpstart_private_hub_model.py | 1 + 2 files changed, 4 insertions(+) diff --git a/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py b/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py index cee5e71aa3..d512915343 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py @@ -148,6 +148,7 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references): model_id=model_id, hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}], + instance_type="ml.g5.2xlarge", ) estimator.fit( @@ -162,6 +163,7 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references): tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}], role=get_sm_session().get_caller_identity_arn(), sagemaker_session=get_sm_session(), + instance_type="ml.g5.2xlarge", ) payload = { @@ -182,6 +184,7 @@ def test_jumpstart_hub_gated_estimator_without_eula(setup, add_model_references) model_id=model_id, hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}], + instance_type="ml.g5.2xlarge", ) with pytest.raises(Exception): estimator.fit( diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 76334330f5..3956c2240d 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -116,6 +116,7 @@ def test_jumpstart_hub_gated_model(setup, add_model_references): predictor = model.deploy( accept_eula=True, tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}], + instance_type="ml.g5.2xlarge", ) payload = model.retrieve_example_payload() From 9be072b4e59f77f321f14db04dd191de665e0f1d Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Fri, 25 Sep 2026 12:30:50 -0700 Subject: [PATCH 3/4] 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. --- src/sagemaker/jumpstart/constants.py | 80 ++++++++++++++--- .../sagemaker/jumpstart/test_constants.py | 86 +++++++++++++++++++ 2 files changed, 155 insertions(+), 11 deletions(-) create mode 100644 tests/unit/sagemaker/jumpstart/test_constants.py diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index b81f97ce3a..45185f9a06 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/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,73 @@ 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 __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 = [ diff --git a/tests/unit/sagemaker/jumpstart/test_constants.py b/tests/unit/sagemaker/jumpstart/test_constants.py new file mode 100644 index 0000000000..12106c0e3d --- /dev/null +++ b/tests/unit/sagemaker/jumpstart/test_constants.py @@ -0,0 +1,86 @@ +# 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(): + """Ensure each test starts and ends with an unresolved proxy cache.""" + _LazyJumpStartSagemakerSession._resolved = False + _LazyJumpStartSagemakerSession._session = None + yield + _LazyJumpStartSagemakerSession._resolved = False + _LazyJumpStartSagemakerSession._session = None + + +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 From d16e4ede91e2541ced796e19f8964904d220c9f7 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Sun, 27 Sep 2026 16:54:43 -0700 Subject: [PATCH 4/4] fix: forward __delattr__ on the lazy JumpStart session proxy _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). --- src/sagemaker/jumpstart/constants.py | 9 ++++ .../sagemaker/jumpstart/test_constants.py | 46 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index 45185f9a06..043e8ddd07 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/src/sagemaker/jumpstart/constants.py @@ -236,6 +236,15 @@ def __getattr__(self, 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()) diff --git a/tests/unit/sagemaker/jumpstart/test_constants.py b/tests/unit/sagemaker/jumpstart/test_constants.py index 12106c0e3d..41e87080a3 100644 --- a/tests/unit/sagemaker/jumpstart/test_constants.py +++ b/tests/unit/sagemaker/jumpstart/test_constants.py @@ -84,3 +84,49 @@ def test_failed_build_degrades_to_none_contract(): 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"