Merge pull request #3 from ShiboXing/add-pt1.13.1-training

Add pt1.13.1 training

Author: Mike Schneider

src/sagemaker/image_uri_config/pytorch.json

@@ -264,6 +264,7 @@
             },
             "1.4.0": {
                 "py_versions": [
+                    "py2",
                     "py3"
                 ],
                 "registries": {

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/sagemaker/image_uris/test_dlc_frameworks.py

@@ -18,6 +18,7 @@
 from tests.unit.sagemaker.image_uris import expected_uris
 
 INSTANCE_TYPES_AND_PROCESSORS = (("ml.c4.xlarge", "cpu"), ("ml.p2.xlarge", "gpu"))
+RENEWED_PYTORCH_INSTANCE_TYPES_AND_PROCESSORS = (("ml.c4.xlarge", "cpu"), ("ml.g4dn.xlarge", "gpu"))
 REGION = "us-west-2"
 
 DLC_ACCOUNT = "763104351884"
@@ -70,7 +71,12 @@ def _test_image_uris(
         "image_scope": scope,
     }
 
-    for instance_type, processor in INSTANCE_TYPES_AND_PROCESSORS:
+    TYPES_AND_PROCESSORS = INSTANCE_TYPES_AND_PROCESSORS
+    if framework == "pytorch" and Version(fw_version) >= Version("1.13"):
+        '''Handle P2 deprecation'''
+        TYPES_AND_PROCESSORS = RENEWED_PYTORCH_INSTANCE_TYPES_AND_PROCESSORS
+
+    for instance_type, processor in TYPES_AND_PROCESSORS:
         uri = image_uris.retrieve(region=REGION, instance_type=instance_type, **base_args)
 
         expected = expected_fn(processor=processor, **expected_fn_args)

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,
     )