change: revise unit test

Author: wanyixia

src/sagemaker/model.py

@@ -652,7 +652,7 @@ def compile(
         self.model_data = job_status["ModelArtifacts"]["S3ModelArtifacts"]
         if target_instance_family is not None:
             if target_instance_family == "ml_eia2":
-                LOGGER.info("You are using target device ml_eia2...")
+                pass
             elif target_instance_family.startswith("ml_"):
                 self.image_uri = self._compilation_image_uri(
                     self.sagemaker_session.boto_region_name,

src/sagemaker/tensorflow/model.py

@@ -289,7 +289,12 @@ def deploy(
 
     def _eia_supported(self):
         """Return true if TF version is EIA enabled"""
-        return [int(s) for s in self.framework_version.split(".")][:2] <= self.LATEST_EIA_VERSION
+        framework_version = [int(s) for s in self.framework_version.split(".")][:2]
+        return (
+            framework_version != [2, 1]
+            and framework_version != [2, 2]
+            and framework_version <= self.LATEST_EIA_VERSION
+        )
 
     def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Prepare the container definition.

tests/unit/sagemaker/tensorflow/test_tfs.py

@@ -142,6 +142,15 @@ def test_tfs_model_image_accelerator_not_supported(sagemaker_session):
 
     model.deploy(instance_type="ml.c4.xlarge", initial_instance_count=1)
 
+    with pytest.raises(AttributeError) as e:
+        model.deploy(
+            instance_type="ml.c4.xlarge",
+            accelerator_type="ml.eia1.medium",
+            initial_instance_count=1,
+        )
+
+    assert str(e.value) == "The TensorFlow version 2.1 doesn't support EIA."
+
 
 def test_tfs_model_with_log_level(sagemaker_session, tensorflow_inference_version):
     model = TensorFlowModel(