From 18f22415148bce747b3f06162e2e1ffd57008617 Mon Sep 17 00:00:00 2001
From: Mohamed Zeidan <zeidmo@amazon.com>
Date: Thu, 30 Apr 2026 18:58:16 -0700
Subject: [PATCH 1/2] 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*p<M{AsM&N^5kwG?qQXdBxJuhl)6R{VsiJSC3k4TG
zfDa{}j0+JPLe5F@pa09{^XvTsKnHaT8DtIQY~-=PP@0H)!Cm2x`Q>0DhY7>Nq4HI7
z#E@;Z`z8u74Oli7VgHdkk8vVBQUx<As`gSN74WM_`H7Q6Qh410mKhpn(shHzO%S=G
zy+(4U;s*6|;_f@4h^9>oC9D`I+o)ibp)$vb!3edl>l3Q=!%7cLAmxv{?Y@B;LuVLx
zTzKLs@Wn&SV=<PVn97L98pIPBCp<lkrl;`ZFl5M`1S82%>?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<QYc2(Rw|US7a?0Hj5Q+bn9-1Z8HVhogk;}t43fy2eP1G!eJlHx
zVK7;<55qsxUAT4c{r~4b^L)?q%`@km@0|0U^PTgK7MP5j8UUaI0KiX^i~&dJIDj0W
zuB0o^ucC2NaK9A*&?2y&22fE^{lt2jAoeG#y1a(UO(h*&0ri{g?FG~gLZ@$jkq&qJ
z_(nk>*SvJNi5f+|mFyAIedTzCd{2gWCC9aikhx{hXVZ#6&7NPABy%mc?|%R=SjHAJ
zh<V*^1S_vh((wT!Lke%(iK4DOc4<Bp2Vr&Yw+4l@ihFuVA?|Z;poZvt176%uFYh+x
z%dOP{Q&LTIJ&`~Xta?OPhjxc~{nbS#!mrA3TU!TJC=#mU2z9oRcXSl6hC0)%SXeMX
zCnqP(v~pEx;Busiv>@8rymO(INj~<t0}WvzAyIPoB7|RBe2OZmqct-LtlR=}OIA~S
zV3J-MQr1;#sj=(+6>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
zEz6<rlSHJ8rTUi^@eYKviS9w2jwkG%mYIb2w_{1HJruyAy1hAgb)AJ}Alt63&W30Q
zEJM&rO}g2>WmHW5+vt_&<&PBLp*}At@$8y~cE&<YqZ=Nts<yLdMmd1ht;|bxZFjx)
z-ImfX`id*&_~t+_VpE-Fuvq~MqWbsPr*O>U2KS1Vj2HU%Y<vgOyy27=i#vLEyo(-k
z$OJ_W_V_C@8--7EW8mo|(s}x~wMo@BECu3{);{%&Fi$%z?cG@xt+mlWZT0umr7zB+
zefzgUl9w*pNuNaN3=~U0e1I<cFcig8=Q;1bmhWm99CTWa(R;p5x;j)Yfvw&idXEt*
z*Wr<`cD%Lh4d5d1m5ov-n2UGV)@d3Ve%zS0e2tBitr11$#Ck3a)(C89p-YCo;0-t<
z5*e-F#F`xjvpsgHxxtwbcy%pWS=Guf-E1L&mf>_CRfcGqD4e*hG51|OBcgIa!8L~M
z<b~>f;(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%<U@K@4yjaqesFV=P0sRNdcGx_=tQgAxp8`r)8ZniqO-X=
z_v{(K<eitcuFo5{ePmsf#OgJM)>{5jI~xFcA;uqFV4DJa70mPqo2Ws4GhWd&<+WiA
zM28o};Wpw;=K82p3m3h!piQwLlP=7PNLGKKh*k`|%)$3UBkE1Y)GerTxK!+U4J~Yt
zUvk$TNm$<CyS`r&yB_Q1l-Xlc81LpKDk=DIF*@QQv-ro4xvf2BTUz`s7!4rr+@f0y
z+T20uhx6di-t8+5=PnW|2z~A_hg3b*ZpjkU1dR5X)Ekz1TX*T4efni12r+Y?7@W}3
zCIJAJ3G1JqVTu3u;;aeJHvjV~>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>qjhV<d>KUAw|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*qPHY<HWjBGbOKfVUWNl>qg%S<7y%sJyOC=R
z!*C(Pd{zMtWMQJLz-udShamPoCK0Gr93D4zy)$B2<Bp&~jPFkHOj@MN#8aM*F6NQu
zfV~M$V^lO_xS7FjkZaPqlGtj;?m@NkSdlD$jRM9^V?`fbq|&w_>TMe~gQ_;iY|n?P
zEQfN>eyXiBSn!48#b=sX*lN~FtXG#fvy;aoPE8<G7smGc7_J_)PVnpGin|%;)NyNX
zd`@l1c&k@ZBYZ+L!h0ROV0y0=<0kBohi}QNt4$M&F>!4h0?XQni?<HBYE(KE3@L#;
z8cvm1X11}|4LqSdCbHwDK9C@@ycYAUtSRYq?d~buxI3hDVsv>xU<tChxr4^pOn74I
zbUI3Al3BgHpV^tupgTS|Xy{Eq@P$)r0ljbXG#Dsf4#(}<ZLLeR$4bP<#dAf^IhI=c
zpPbEH<YUoo0$Nrjt?Bia$LMTda3Zy&?F#J@Q4_g^5do_5ZSt@)$ht(j`AmVHi>bU)
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~PSvPHXY<b-;#
zc!^K$LE&V*{0=4l1}2n*Rg=|DY~Ra4bE@n^@<gj-{+cs5eMP3_0ppvLFTM)M0=BIJ
zK^2OVC=)Ue`7dO{2KE7qosdSSigQ<gT}{{vogf`vF3f0ep*56}&VjeA2eIks@u7C3
zc;-5sHP)&++T#0o(k2XRsHJ+#1Wl|b*2GW-7L__;7?85z=*8i6>rl>TBeZel$H1>P
zn<;JDWBBS!+Kh)?1SSWflO0Hh#s~RDqXm+djpn}%zE%Khc*aglGqrJ6@OY@@C#5d{
zSF(zr`#Z8B_^<YC<<}9NJF)p!N7j-aP%Gsmps~go7j##|_je{@nM-oAW_dKl-Y7x@
zX<TQ$M|KuxM^}zxeVMof91>DWhs?|&4*SnU1p7tj8Fj*fH8VZwk!>D$AN8hcxM9y~
zO7zGtEkOE@AK$GcF+hZERqPw#BodW1RTU>BsLJ>EN8l{2&ido$?|3JqcSfX1*-&fN
zS|xp*PsJUehEJXMCvBUL5?mP#6f2L%&)gD3<F867Ob=QnWF=cMOL?o9V}+4mOzH8+
z*>tN$Cbr9n^}g$K%kGhO=tq|@c>}CtO%VRs&GN}RrH%(C%~NZ22yy<Ms&g>BIUS!i
z@HwBp{?&WQ6<lBG2g+7Dxe{DwVPzP|ZWZ?;pA;{-7<hK8g95jFVLUjHG7xcUO{C9n
z{}W#(@6LiJrhTRtktFFc^5}Xest)Vj?ADRFHZF-9D8_s3b>Wg$`)bm1XI~<)FG@4=
z#+$l0EH#t%mR_Sf-uGKw;=cWkN<d8YVDXD9`7aApr7t3ao{_jIg;95BdEew_piqoe
z$TL*ZDz*|ebD%hz0Qe^1{D3<c)?iM-#w1d>TfpQ+mH+UJ9_>_r{nJ!5qo_qepwV0M
z;Bj9TkBDa~qv;=C#={zJln*wt`?m60FwtM3F$_s<zU-_~azohhyc5L~uU-$DLqhg$
z>R?pUP5DcH_H@p}v+@PpPuZ_#A9F7%eArm=P>Y%+p=(I~`uBNQtpvx<NC^M{X^!^P
zALn6OLQl1W+FMx=rbYo5cb7Hw_Vzy4P8E<q>22lCuFv3h^_<o&Fi5bUPeV<!L%m&t
zzn72Lm-_is_yu~_{_Is0K@!ZT9Z9|)?Ydu0Cr?mJTv}ifQfk0IgJ%NE5&9D!{|2Q=
zNKXR}$JnEAl{gC=5+7mvr|1!AO+5S$4Mfk6Vgw?rJtQ<?`}5380J%r)|1o$(FmXuv
zWPd=5Lyv!yAcBiS>iu2fDAf38St8gtr2U^{k3x@s)<OgyhjfMD0uc@!T5zc44=8e|
z>&JoY{GlSk&hfwM`W>1a-U;#Fl%vxx5u6-S!T+Z4=WF{8RerqtqcbLPmOdnF0%H94
zAbdCc=W~by=+R?{@i|&{e@#+bLNWf`_<B?WF<*{W53xKC=^Y_`juxvH_!z--0DzY8
NF(%}L)JfvE{{T!qvzh<^

literal 0
HcmV?d00001

diff --git a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py
index 8a19baaf40..7341d3ad6c 100644
--- a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py
+++ b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py
@@ -95,9 +95,10 @@ def test_prepacked_jumpstart_model(setup):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.p3.2xlarge instance
+    # uses ml.g4dn.xlarge instance
     predictor = model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g4dn.xlarge",
     )
 
     response = predictor.predict("hello world!")
@@ -120,7 +121,7 @@ def test_model_package_arn_jumpstart_model(setup):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.g5.2xlarge instance
+    # uses ml.g4dn.2xlarge instance
     predictor = model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
     )
@@ -175,7 +176,7 @@ def test_jumpstart_gated_model(setup):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.g5.2xlarge instance
+    # uses ml.g4dn.2xlarge instance
     predictor = model.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         accept_eula=True,
@@ -206,6 +207,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup):
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         accept_eula=True,
         endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED,
+        instance_type="ml.g5.2xlarge",
     )
 
     predictor = retrieve_default(
@@ -262,7 +264,7 @@ def test_jumpstart_model_register(setup):
 
     # uses  instance
     predictor = model_package.deploy(
-        instance_type="ml.p3.2xlarge",
+        instance_type="ml.g4dn.xlarge",
         initial_instance_count=1,
     )
 
@@ -379,7 +381,7 @@ def test_jumpstart_model_with_deployment_configs(setup):
 
     model.set_deployment_config(
         configs[0]["ConfigName"],
-        "ml.g5.2xlarge",
+        "ml.g4dn.2xlarge",
     )
     assert model.config_name == configs[0]["ConfigName"]
 
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 a6e33f1bdf..cee5e71aa3 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
@@ -69,6 +69,7 @@ def test_jumpstart_hub_estimator(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.g4dn.xlarge",
     )
 
     estimator.fit(
@@ -85,9 +86,10 @@ def test_jumpstart_hub_estimator(setup, add_model_references):
         model_version=model_version,
     )
 
-    # uses ml.p3.2xlarge instance
+    # uses ml.g4dn.xlarge instance
     predictor = estimator.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        instance_type="ml.g4dn.xlarge",
     )
 
     response = predictor.predict(["hello", "world"])
@@ -107,6 +109,7 @@ def test_jumpstart_hub_estimator_with_session(setup, add_model_references):
         sagemaker_session=sagemaker_session,
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME],
+        instance_type="ml.g4dn.xlarge",
     )
 
     estimator.fit(
@@ -124,11 +127,12 @@ def test_jumpstart_hub_estimator_with_session(setup, add_model_references):
         sagemaker_session=get_sm_session(),
     )
 
-    # uses ml.p3.2xlarge instance
+    # uses ml.g4dn.xlarge instance
     predictor = estimator.deploy(
         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.g4dn.xlarge",
     )
 
     response = predictor.predict(["hello", "world"])
diff --git a/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py b/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py
index 7cb0f34fbf..429a94f8db 100644
--- a/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py
+++ b/tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py
@@ -34,8 +34,8 @@
 def test_jumpstart_transfer_learning_retrieve_functions(setup):
 
     model_id, model_version = "huggingface-spc-bert-base-cased", "2.0.3"
-    training_instance_type = "ml.p3.2xlarge"
-    inference_instance_type = "ml.p2.xlarge"
+    training_instance_type = "ml.g4dn.xlarge"
+    inference_instance_type = "ml.g4dn.xlarge"
 
     # training
     print("Starting training...")
diff --git a/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py b/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py
index 329b069fee..5b4c89c5b9 100644
--- a/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py
+++ b/tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py
@@ -53,12 +53,8 @@ def test_jumpstart_transfer_learning_estimator_class(setup):
 
     model_id, model_version = "huggingface-spc-bert-base-cased", "1.2.3"
 
-    inference_instance_type = instance_types.retrieve_default(
-        model_id=model_id, model_version=model_version, scope="inference"
-    )
-    training_instance_type = instance_types.retrieve_default(
-        model_id=model_id, model_version=model_version, scope="training"
-    )
+    inference_instance_type = "ml.g4dn.xlarge"
+    training_instance_type = "ml.g4dn.xlarge"
     instance_count = 1
 
     print("Starting training...")
diff --git a/tests/integ/test_training_compiler.py b/tests/integ/test_training_compiler.py
index 1251eb0723..f4fe7d293e 100644
--- a/tests/integ/test_training_compiler.py
+++ b/tests/integ/test_training_compiler.py
@@ -30,7 +30,7 @@
 
 @pytest.fixture(scope="module")
 def gpu_instance_type(request):
-    return "ml.p3.2xlarge"
+    return "ml.g4dn.xlarge"
 
 
 @pytest.fixture(scope="module")
@@ -86,8 +86,8 @@ def skip_if_incompatible(gpu_instance_type, request):
 @pytest.mark.parametrize(
     "gpu_instance_type,instance_count",
     [
-        pytest.param("ml.p3.2xlarge", 1, marks=pytest.mark.release),
-        pytest.param("ml.p3.16xlarge", 2),
+        pytest.param("ml.g4dn.xlarge", 1, marks=pytest.mark.release),
+        pytest.param("ml.g4dn.12xlarge", 2),
     ],
 )
 @pytest.mark.skipif(
@@ -138,8 +138,8 @@ def test_huggingface_pytorch(
 @pytest.mark.parametrize(
     "gpu_instance_type,instance_count",
     [
-        pytest.param("ml.p3.2xlarge", 1, marks=pytest.mark.release),
-        pytest.param("ml.p3.16xlarge", 2),
+        pytest.param("ml.g4dn.xlarge", 1, marks=pytest.mark.release),
+        pytest.param("ml.g4dn.12xlarge", 2),
     ],
 )
 @pytest.mark.skip("Temporarily skip to unblock")

From 7412b08e10b90b9c31e6a3848756946b5181bca7 Mon Sep 17 00:00:00 2001
From: Mohamed Zeidan <zeidmo@amazon.com>
Date: Fri, 1 May 2026 00:30:50 -0700
Subject: [PATCH 2/2] 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()