fix: p2 error in unit test with a fixture

Author: Shibo Xing

tests/conftest.py

@@ -86,6 +86,7 @@
     "huggingface_training_compiler",
 )
 
+PYTORCH_RENEWED_GPU = "ml.g4dn.xlarge"
 
 def pytest_addoption(parser):
     parser.addoption("--sagemaker-client-config", action="store", default=None)
@@ -514,6 +515,21 @@ def gpu_instance_type(sagemaker_session, request):
     else:
         return "ml.p3.2xlarge"
 
+@pytest.fixture()
+def gpu_pytorch_instance_type(sagemaker_session, request):
+    if "pytorch_inference_version" in request.fixturenames:
+        fw_version = request.getfixturevalue("pytorch_inference_version")
+    else:
+        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"
 
 @pytest.fixture(scope="session")
 def gpu_instance_type_list(sagemaker_session, request):

tests/unit/test_pytorch.py

@@ -302,7 +302,7 @@ def test_create_model_with_custom_image(name_from_base, sagemaker_session):
 @patch("sagemaker.estimator.name_from_base", return_value=JOB_NAME)
 @patch("time.time", return_value=TIME)
 def test_pytorch(
-    time, name_from_base, sagemaker_session, pytorch_inference_version, pytorch_inference_py_version
+    time, name_from_base, sagemaker_session, pytorch_inference_version, pytorch_inference_py_version, gpu_pytorch_instance_type
 ):
     pytorch = PyTorch(
         entry_point=SCRIPT_PATH,
@@ -339,24 +339,24 @@ def test_pytorch(
         REGION,
         version=pytorch_inference_version,
         py_version=pytorch_inference_py_version,
-        instance_type=GPU,
+        instance_type=gpu_pytorch_instance_type,
         image_scope="inference",
     )
 
-    actual_environment = model.prepare_container_def(GPU)
+    actual_environment = model.prepare_container_def(gpu_pytorch_instance_type)
     submit_directory = actual_environment["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"]
     model_url = actual_environment["ModelDataUrl"]
     expected_environment = _get_environment(submit_directory, model_url, expected_image_uri)
     assert actual_environment == expected_environment
 
     assert "cpu" in model.prepare_container_def(CPU)["Image"]
-    predictor = pytorch.deploy(1, GPU)
+    predictor = pytorch.deploy(1, gpu_pytorch_instance_type)
     assert isinstance(predictor, PyTorchPredictor)
 
 
 @patch("sagemaker.utils.repack_model", MagicMock())
 @patch("sagemaker.utils.create_tar_file", MagicMock())
-def test_model(sagemaker_session, pytorch_inference_version, pytorch_inference_py_version):
+def test_model(sagemaker_session, pytorch_inference_version, pytorch_inference_py_version, gpu_pytorch_instance_type):
     model = PyTorchModel(
         MODEL_DATA,
         role=ROLE,
@@ -365,21 +365,22 @@ def test_model(sagemaker_session, pytorch_inference_version, pytorch_inference_p
         py_version=pytorch_inference_py_version,
         sagemaker_session=sagemaker_session,
     )
-    predictor = model.deploy(1, GPU)
+    predictor = model.deploy(1, gpu_pytorch_instance_type)
     assert isinstance(predictor, PyTorchPredictor)
 
 
 @patch("sagemaker.utils.create_tar_file", MagicMock())
 @patch("sagemaker.utils.repack_model")
-def test_mms_model(repack_model, sagemaker_session):
+@pytest.mark.parametrize("gpu_pytorch_instance_type", ["1.2"], indirect=True)
+def test_mms_model(repack_model, sagemaker_session, gpu_pytorch_instance_type):
     PyTorchModel(
         MODEL_DATA,
         role=ROLE,
         entry_point=SCRIPT_PATH,
         sagemaker_session=sagemaker_session,
         framework_version="1.2",
         py_version="py3",
-    ).deploy(1, GPU)
+    ).deploy(1, gpu_pytorch_instance_type)
 
     repack_model.assert_called_with(
         dependencies=[],
@@ -428,6 +429,7 @@ def test_model_custom_serialization(
     sagemaker_session,
     pytorch_inference_version,
     pytorch_inference_py_version,
+    gpu_pytorch_instance_type
 ):
     model = PyTorchModel(
         MODEL_DATA,
@@ -441,7 +443,7 @@ def test_model_custom_serialization(
     custom_deserializer = Mock()
     predictor = model.deploy(
         1,
-        GPU,
+        gpu_pytorch_instance_type,
         serializer=custom_serializer,
         deserializer=custom_deserializer,
     )