check framework and py version in accelerator validate method

Author: Deng

src/sagemaker/fw_utils.py

@@ -60,6 +60,7 @@
     "mxnet-serving",
     "pytorch-serving",
 ]
+PY2_RESTRICTED_EIA_FRAMEWORKS = ["pytorch-serving"]
 VALID_ACCOUNTS_BY_REGION = {"us-gov-west-1": "246785580436", "us-iso-east-1": "744548109606"}
 ASIMOV_VALID_ACCOUNTS_BY_REGION = {"us-iso-east-1": "886529160074"}
 OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "057415533634", "me-south-1": "724002660598"}
@@ -215,6 +216,7 @@ def create_image_uri(
 
     if _accelerator_type_valid_for_framework(
         framework=framework,
+        py_version=py_version,
         accelerator_type=accelerator_type,
         optimized_families=optimized_families,
     ):
@@ -267,17 +269,23 @@ def create_image_uri(
 
 
 def _accelerator_type_valid_for_framework(
-    framework, accelerator_type=None, optimized_families=None
+    framework, py_version, accelerator_type=None, optimized_families=None
 ):
     """
     Args:
         framework:
+        py_version:
         accelerator_type:
         optimized_families:
     """
     if accelerator_type is None:
         return False
 
+    if py_version == "py2" and framework in PY2_RESTRICTED_EIA_FRAMEWORKS:
+        raise ValueError(
+            "{} is not supported with Amazon Elastic Inference in Python 2.".format(framework)
+        )
+
     if framework not in VALID_EIA_FRAMEWORKS:
         raise ValueError(
             "{} is not supported with Amazon Elastic Inference. Currently only "

src/sagemaker/pytorch/model.py

@@ -142,9 +142,6 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             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

tests/unit/test_fw_utils.py

@@ -328,6 +328,20 @@ def test_pytorch_eia_images():
     )
 
 
+def test_pytorch_eia_py2_error():
+    error_message = "pytorch-serving is not supported with Amazon Elastic Inference in Python 2."
+    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_pytorch.py

@@ -355,7 +355,9 @@ def test_model_image_accelerator(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)
+    assert "pytorch-serving is not supported with Amazon Elastic Inference in Python 2." in str(
+        error
+    )
 
 
 def test_train_image_default(sagemaker_session):