address comments: accelerator_type, script name, framework and python version check

Author: Deng

doc/using_pytorch.rst

@@ -259,7 +259,7 @@ In order to attach an Elastic Inference accelerator to your endpoint provide the
 
   predictor = pytorch_estimator.deploy(instance_type='ml.m4.xlarge',
                                        initial_instance_count=1,
-                                       accelerator_type='ml.eia1.medium')
+                                       accelerator_type='ml.eia2.medium')
 
 The SageMaker PyTorch Model Server
 ==================================

src/sagemaker/fw_utils.py

@@ -58,7 +58,6 @@
     "tensorflow-serving",
     "mxnet",
     "mxnet-serving",
-    "pytorch",
     "pytorch-serving",
 ]
 VALID_ACCOUNTS_BY_REGION = {"us-gov-west-1": "246785580436", "us-iso-east-1": "744548109606"}
@@ -90,7 +89,7 @@
     "mxnet-serving-eia": [1, 4, 1],
     "pytorch": [1, 2, 0],
     "pytorch-serving": [1, 2, 0],
-    "pytorch-serving-eia": {"py3": [1, 3, 1]},
+    "pytorch-serving-eia": [1, 3, 1],
 }
 
 DEBUGGER_UNSUPPORTED_REGIONS = ["us-gov-west-1", "us-iso-east-1"]
@@ -126,8 +125,6 @@ def _is_dlc_version(framework, framework_version, py_version):
     """
     lowest_version_list = MERGED_FRAMEWORKS_LOWEST_VERSIONS.get(framework)
     if isinstance(lowest_version_list, dict):
-        if py_version not in lowest_version_list:
-            raise ValueError("{} is not supported in {}.".format(framework, py_version))
         lowest_version_list = lowest_version_list[py_version]
 
     if lowest_version_list:

src/sagemaker/pytorch/model.py

@@ -136,12 +136,15 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
                 For example, 'ml.p2.xlarge'.
             accelerator_type (str): The Elastic Inference accelerator type to
                 deploy to the instance for loading and making inferences to the
-                model. Currently supported with PyTorch 1.3.1 Python 3.
+                model.
 
         Returns:
             dict[str, str]: A container definition object usable with the
             CreateModel API.
         """
+        if accelerator_type and self.py_version == "py2":
+            raise ValueError("PyTorch EIA is not supported in Python 2.")
+
         deploy_image = self.image
         if not deploy_image:
             region_name = self.sagemaker_session.boto_session.region_name
@@ -169,7 +172,7 @@ def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
                 (cpu/gpu/family-specific optimized).
             accelerator_type (str): The Elastic Inference accelerator type to
                 deploy to the instance for loading and making inferences to the
-                model. Currently supported with PyTorch 1.3.1 Python 3.
+                model.
 
         Returns:
             str: The appropriate image URI based on the given parameters.

tests/data/pytorch_eia/mnist.py renamed to tests/data/pytorch_eia/empty_inference_script.py

@@ -29,7 +29,7 @@
 
 EIA_DIR = os.path.join(DATA_DIR, "pytorch_eia")
 EIA_MODEL = os.path.join(EIA_DIR, "model_mnist.tar.gz")
-EIA_SCRIPT = os.path.join(EIA_DIR, "mnist.py")
+EIA_SCRIPT = os.path.join(EIA_DIR, "empty_inference_script.py")
 
 
 @pytest.fixture(scope="module", name="pytorch_training_job")

tests/integ/test_pytorch_train.py

@@ -328,20 +328,6 @@ def test_pytorch_eia_images():
     )
 
 
-def test_pytorch_eia_py2_error():
-    error_message = "pytorch-serving-eia is not supported in py2."
-    with pytest.raises(ValueError) as error:
-        fw_utils.create_image_uri(
-            "us-east-1",
-            "pytorch-serving",
-            "ml.c4.2xlarge",
-            "1.3.1",
-            "py2",
-            accelerator_type="ml.eia1.large",
-        )
-    assert error_message in str(error)
-
-
 def test_create_image_uri_override_account():
     image_uri = fw_utils.create_image_uri(
         "us-west-1", MOCK_FRAMEWORK, "ml.p3.2xlarge", "1.0rc", "py3", account="fake"

tests/unit/test_fw_utils.py

@@ -352,6 +352,20 @@ def test_model_image_accelerator(sagemaker_session):
     assert isinstance(predictor, PyTorchPredictor)
 
 
+@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
+def test_model_image_accelerator(sagemaker_session):
+    with pytest.raises(ValueError) as error:
+        model = PyTorchModel(
+            MODEL_DATA,
+            role=ROLE,
+            entry_point=SCRIPT_PATH,
+            sagemaker_session=sagemaker_session,
+            py_version="py2",
+        )
+        model.deploy(1, CPU, accelerator_type=ACCELERATOR_TYPE)
+    assert "PyTorch EIA is not supported in Python 2." in str(error)
+
+
 def test_train_image_default(sagemaker_session):
     pytorch = PyTorch(
         entry_point=SCRIPT_PATH,