@@ -67,8 +67,12 @@ class EstimatorTest(unittest.TestCase):
67
67
@mock .patch ("sagemaker.jumpstart.factory.model.JUMPSTART_LOGGER" )
68
68
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
69
69
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
70
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
71
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
70
+ @mock .patch (
71
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
72
+ )
73
+ @mock .patch (
74
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
75
+ )
72
76
@mock .patch ("sagemaker.jumpstart.utils.validate_model_id_and_get_type" )
73
77
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
74
78
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
@@ -193,8 +197,12 @@ def test_non_prepacked(
193
197
)
194
198
195
199
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
196
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
197
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
200
+ @mock .patch (
201
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
202
+ )
203
+ @mock .patch (
204
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
205
+ )
198
206
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
199
207
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
200
208
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -293,8 +301,12 @@ def test_prepacked(
293
301
294
302
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
295
303
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
296
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
297
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
304
+ @mock .patch (
305
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
306
+ )
307
+ @mock .patch (
308
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
309
+ )
298
310
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
299
311
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
300
312
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -411,8 +423,12 @@ def test_gated_model_s3_uri(
411
423
)
412
424
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
413
425
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
414
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
415
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
426
+ @mock .patch (
427
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
428
+ )
429
+ @mock .patch (
430
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
431
+ )
416
432
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
417
433
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
418
434
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -559,8 +575,12 @@ def test_gated_model_non_model_package_s3_uri(
559
575
560
576
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
561
577
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
562
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
563
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
578
+ @mock .patch (
579
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
580
+ )
581
+ @mock .patch (
582
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
583
+ )
564
584
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
565
585
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
566
586
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -751,8 +771,12 @@ def test_estimator_use_kwargs(self):
751
771
@mock .patch ("sagemaker.jumpstart.factory.estimator.environment_variables.retrieve_default" )
752
772
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
753
773
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
754
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
755
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
774
+ @mock .patch (
775
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
776
+ )
777
+ @mock .patch (
778
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
779
+ )
756
780
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
757
781
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
758
782
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1160,8 +1184,12 @@ def test_validate_model_id_and_get_type(
1160
1184
1161
1185
@mock .patch ("sagemaker.jumpstart.estimator.get_default_predictor" )
1162
1186
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1163
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1164
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1187
+ @mock .patch (
1188
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1189
+ )
1190
+ @mock .patch (
1191
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1192
+ )
1165
1193
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1166
1194
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1167
1195
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1219,8 +1247,12 @@ def test_no_predictor_returns_default_predictor(
1219
1247
1220
1248
@mock .patch ("sagemaker.jumpstart.estimator.get_default_predictor" )
1221
1249
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1222
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1223
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1250
+ @mock .patch (
1251
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1252
+ )
1253
+ @mock .patch (
1254
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1255
+ )
1224
1256
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1225
1257
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1226
1258
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1269,8 +1301,12 @@ def test_no_predictor_yes_async_inference_config(
1269
1301
1270
1302
@mock .patch ("sagemaker.jumpstart.estimator.get_default_predictor" )
1271
1303
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1272
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1273
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1304
+ @mock .patch (
1305
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1306
+ )
1307
+ @mock .patch (
1308
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1309
+ )
1274
1310
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1275
1311
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1276
1312
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1321,8 +1357,12 @@ def test_yes_predictor_returns_unmodified_predictor(
1321
1357
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1322
1358
@mock .patch ("sagemaker.jumpstart.factory.estimator._model_supports_incremental_training" )
1323
1359
@mock .patch ("sagemaker.jumpstart.factory.estimator.JUMPSTART_LOGGER.warning" )
1324
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1325
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1360
+ @mock .patch (
1361
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1362
+ )
1363
+ @mock .patch (
1364
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1365
+ )
1326
1366
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1327
1367
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1328
1368
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1375,8 +1415,12 @@ def test_incremental_training_with_unsupported_model_logs_warning(
1375
1415
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1376
1416
@mock .patch ("sagemaker.jumpstart.factory.estimator._model_supports_incremental_training" )
1377
1417
@mock .patch ("sagemaker.jumpstart.factory.estimator.JUMPSTART_LOGGER.warning" )
1378
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1379
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1418
+ @mock .patch (
1419
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1420
+ )
1421
+ @mock .patch (
1422
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1423
+ )
1380
1424
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1381
1425
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1382
1426
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1425,8 +1469,12 @@ def test_incremental_training_with_supported_model_doesnt_log_warning(
1425
1469
1426
1470
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
1427
1471
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1428
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1429
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1472
+ @mock .patch (
1473
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1474
+ )
1475
+ @mock .patch (
1476
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1477
+ )
1430
1478
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1431
1479
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1432
1480
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1486,8 +1534,12 @@ def test_estimator_sets_different_inference_instance_depending_on_training_insta
1486
1534
1487
1535
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
1488
1536
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1489
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1490
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1537
+ @mock .patch (
1538
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
1539
+ )
1540
+ @mock .patch (
1541
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1542
+ )
1491
1543
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1492
1544
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1493
1545
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.fit" )
@@ -1564,10 +1616,11 @@ def test_training_passes_role_to_deploy(
1564
1616
@mock .patch ("sagemaker.utils.sagemaker_timestamp" )
1565
1617
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1566
1618
@mock .patch (
1567
- "sagemaker.jumpstart.factory.model.DEFAULT_JUMPSTART_SAGEMAKER_SESSION" , sagemaker_session
1619
+ "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix" ,
1620
+ sagemaker_session ,
1568
1621
)
1569
1622
@mock .patch (
1570
- "sagemaker.jumpstart.factory.estimator.DEFAULT_JUMPSTART_SAGEMAKER_SESSION " ,
1623
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix " ,
1571
1624
sagemaker_session ,
1572
1625
)
1573
1626
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
@@ -1645,7 +1698,9 @@ def test_training_passes_session_to_deploy(
1645
1698
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.deploy" )
1646
1699
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1647
1700
@mock .patch ("sagemaker.jumpstart.factory.estimator._retrieve_estimator_init_kwargs" )
1648
- @mock .patch ("sagemaker.jumpstart.factory.estimator.Session" )
1701
+ @mock .patch (
1702
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1703
+ )
1649
1704
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1650
1705
@mock .patch ("sagemaker.jumpstart.estimator.JumpStartModelsAccessor.reset_cache" )
1651
1706
@mock .patch ("sagemaker.jumpstart.factory.estimator.JUMPSTART_DEFAULT_REGION_NAME" , region )
@@ -1725,7 +1780,9 @@ def test_model_id_not_found_refeshes_cache_training(
1725
1780
)
1726
1781
1727
1782
@mock .patch ("sagemaker.jumpstart.estimator.validate_model_id_and_get_type" )
1728
- @mock .patch ("sagemaker.jumpstart.factory.model.Session" )
1783
+ @mock .patch (
1784
+ "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix"
1785
+ )
1729
1786
@mock .patch ("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs" )
1730
1787
@mock .patch ("sagemaker.jumpstart.estimator.Estimator.__init__" )
1731
1788
@mock .patch ("sagemaker.jumpstart.factory.estimator.JUMPSTART_DEFAULT_REGION_NAME" , region )
0 commit comments