Skip to content

Commit d5c5474

Browse files
samrudssamruds
committed
fix formatting
1 parent fd0e17d commit d5c5474

File tree

1 file changed

+56
-38
lines changed

1 file changed

+56
-38
lines changed

‎tests/unit/sagemaker/serve/builder/test_model_builder.py

Lines changed: 56 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1247,7 +1247,7 @@ def test_build_for_transformers_happy_case_with_values(
12471247
mock_build_for_transformers.assert_called_once()
12481248

12491249
@patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_djl", Mock())
1250-
@patch("sagemaker.serve.utils.hardware_detector._get_gpu_info")
1250+
@patch("sagemaker.serve.builder.model_builder._get_gpu_info")
12511251
@patch("sagemaker.serve.builder.model_builder.ModelBuilder._total_inference_model_size_mib")
12521252
@patch("sagemaker.image_uris.retrieve")
12531253
@patch("sagemaker.djl_inference.model.urllib")
@@ -1257,16 +1257,16 @@ def test_build_for_transformers_happy_case_with_values(
12571257
@patch("sagemaker.model_uris.retrieve")
12581258
@patch("sagemaker.serve.builder.model_builder._ServeSettings")
12591259
def test_build_for_transformers_happy_case_with_valid_gpu_info(
1260-
self,
1261-
mock_serveSettings,
1262-
mock_model_uris_retrieve,
1263-
mock_llm_utils_json,
1264-
mock_llm_utils_urllib,
1265-
mock_model_json,
1266-
mock_model_urllib,
1267-
mock_image_uris_retrieve,
1268-
mock_total_inference_model_size_mib,
1269-
mock_try_fetch_gpu_info,
1260+
self,
1261+
mock_serveSettings,
1262+
mock_model_uris_retrieve,
1263+
mock_llm_utils_json,
1264+
mock_llm_utils_urllib,
1265+
mock_model_json,
1266+
mock_model_urllib,
1267+
mock_image_uris_retrieve,
1268+
mock_total_inference_model_size_mib,
1269+
mock_try_fetch_gpu_info,
12701270
):
12711271
mock_setting_object = mock_serveSettings.return_value
12721272
mock_setting_object.role_arn = mock_role_arn
@@ -1285,7 +1285,9 @@ def test_build_for_transformers_happy_case_with_valid_gpu_info(
12851285

12861286
model_builder = ModelBuilder(model="stable-diffusion")
12871287
model_builder.build(sagemaker_session=mock_session)
1288-
1288+
self.assertEqual(
1289+
model_builder._try_fetch_gpu_info(), INSTANCE_GPU_INFO[1] / INSTANCE_GPU_INFO[0]
1290+
)
12891291
self.assertEqual(model_builder._can_fit_on_single_gpu(), False)
12901292

12911293
@patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers", Mock())
@@ -1300,17 +1302,17 @@ def test_build_for_transformers_happy_case_with_valid_gpu_info(
13001302
@patch("sagemaker.model_uris.retrieve")
13011303
@patch("sagemaker.serve.builder.model_builder._ServeSettings")
13021304
def test_build_for_transformers_happy_case_with_valid_gpu_fallback(
1303-
self,
1304-
mock_serveSettings,
1305-
mock_model_uris_retrieve,
1306-
mock_llm_utils_json,
1307-
mock_llm_utils_urllib,
1308-
mock_model_json,
1309-
mock_model_urllib,
1310-
mock_image_uris_retrieve,
1311-
mock_total_inference_model_size_mib,
1312-
mock_gpu_fallback,
1313-
mock_try_fetch_gpu_info,
1305+
self,
1306+
mock_serveSettings,
1307+
mock_model_uris_retrieve,
1308+
mock_llm_utils_json,
1309+
mock_llm_utils_urllib,
1310+
mock_model_json,
1311+
mock_model_urllib,
1312+
mock_image_uris_retrieve,
1313+
mock_total_inference_model_size_mib,
1314+
mock_gpu_fallback,
1315+
mock_try_fetch_gpu_info,
13141316
):
13151317
mock_setting_object = mock_serveSettings.return_value
13161318
mock_setting_object.role_arn = mock_role_arn
@@ -1324,12 +1326,20 @@ def test_build_for_transformers_happy_case_with_valid_gpu_fallback(
13241326
mock_model_urllib.request.Request.side_effect = Mock()
13251327
mock_try_fetch_gpu_info.side_effect = ValueError
13261328
mock_gpu_fallback.return_value = INSTANCE_GPU_INFO
1327-
mock_total_inference_model_size_mib.return_value = INSTANCE_GPU_INFO[1]/INSTANCE_GPU_INFO[0] - 1
1329+
mock_total_inference_model_size_mib.return_value = (
1330+
INSTANCE_GPU_INFO[1] / INSTANCE_GPU_INFO[0] - 1
1331+
)
13281332

13291333
mock_image_uris_retrieve.return_value = "https://some-image-uri"
13301334

1331-
model_builder = ModelBuilder(model="stable-diffusion", sagemaker_session=mock_session, instance_type=mock_instance_type)
1332-
self.assertEqual(model_builder._try_fetch_gpu_info(), INSTANCE_GPU_INFO[1]/INSTANCE_GPU_INFO[0])
1335+
model_builder = ModelBuilder(
1336+
model="stable-diffusion",
1337+
sagemaker_session=mock_session,
1338+
instance_type=mock_instance_type,
1339+
)
1340+
self.assertEqual(
1341+
model_builder._try_fetch_gpu_info(), INSTANCE_GPU_INFO[1] / INSTANCE_GPU_INFO[0]
1342+
)
13331343
self.assertEqual(model_builder._can_fit_on_single_gpu(), True)
13341344

13351345
@patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers", Mock())
@@ -1343,16 +1353,16 @@ def test_build_for_transformers_happy_case_with_valid_gpu_fallback(
13431353
@patch("sagemaker.model_uris.retrieve")
13441354
@patch("sagemaker.serve.builder.model_builder._ServeSettings")
13451355
def test_build_for_transformers_happy_case_hugging_face_responses(
1346-
self,
1347-
mock_serveSettings,
1348-
mock_model_uris_retrieve,
1349-
mock_llm_utils_json,
1350-
mock_llm_utils_urllib,
1351-
mock_model_json,
1352-
mock_model_urllib,
1353-
mock_image_uris_retrieve,
1354-
mock_gather_data,
1355-
mock_parser,
1356+
self,
1357+
mock_serveSettings,
1358+
mock_model_uris_retrieve,
1359+
mock_llm_utils_json,
1360+
mock_llm_utils_urllib,
1361+
mock_model_json,
1362+
mock_model_urllib,
1363+
mock_image_uris_retrieve,
1364+
mock_gather_data,
1365+
mock_parser,
13561366
):
13571367
mock_setting_object = mock_serveSettings.return_value
13581368
mock_setting_object.role_arn = mock_role_arn
@@ -1370,12 +1380,20 @@ def test_build_for_transformers_happy_case_hugging_face_responses(
13701380
mock_gather_data.return_value = [[1, 1, 1, 1]]
13711381
product = MIB_CONVERSION_FACTOR * 1 * MEMORY_BUFFER_MULTIPLIER
13721382

1373-
model_builder = ModelBuilder(model="stable-diffusion", sagemaker_session=mock_session, instance_type=mock_instance_type)
1383+
model_builder = ModelBuilder(
1384+
model="stable-diffusion",
1385+
sagemaker_session=mock_session,
1386+
instance_type=mock_instance_type,
1387+
)
13741388
self.assertEqual(model_builder._total_inference_model_size_mib(), product)
13751389

13761390
mock_parser.return_value = Mock()
13771391
mock_gather_data.return_value = None
1378-
model_builder = ModelBuilder(model="stable-diffusion", sagemaker_session=mock_session, instance_type=mock_instance_type)
1392+
model_builder = ModelBuilder(
1393+
model="stable-diffusion",
1394+
sagemaker_session=mock_session,
1395+
instance_type=mock_instance_type,
1396+
)
13791397
with self.assertRaises(ValueError) as _:
13801398
model_builder._total_inference_model_size_mib()
13811399

0 commit comments

Comments
 (0)