fixed linting issues

Author: mvsusp

src/sagemaker/fw_utils.py

@@ -64,7 +64,6 @@
     "mxnet": "mxnet-training",
     "tensorflow-serving": "tensorflow-inference",
     "tensorflow-serving-eia": "tensorflow-inference-eia",
-    "mxnet-serving": "mxnet-inference",
     "mxnet-serving-eia": "mxnet-inference-eia",
 }
 
@@ -73,7 +72,6 @@
     "mxnet": [1, 4, 1],
     "tensorflow-serving": [1, 13, 0],
     "tensorflow-serving-eia": [1, 14, 0],
-    "mxnet-serving": [1, 4, 1],
     "mxnet-serving-eia": [1, 4, 1],
 }
 
@@ -212,6 +210,13 @@ def create_image_uri(
     if py_version and py_version not in VALID_PY_VERSIONS:
         raise ValueError("invalid py_version argument: {}".format(py_version))
 
+    if _accelerator_type_valid_for_framework(
+        framework=framework,
+        accelerator_type=accelerator_type,
+        optimized_families=optimized_families,
+    ):
+        framework += "-eia"
+
     # Handle Account Number for Gov Cloud and frameworks with DLC merged images
     account = _registry_id(
         region=region,
@@ -242,13 +247,6 @@ def create_image_uri(
         else:
             device_type = "cpu"
 
-    if _accelerator_type_valid_for_framework(
-        framework=framework,
-        accelerator_type=accelerator_type,
-        optimized_families=optimized_families,
-    ):
-        framework += "-eia"
-
     using_merged_images = _using_merged_images(region, framework, py_version, framework_version)
 
     if not py_version or (using_merged_images and framework == "tensorflow-serving-eia"):

tests/unit/test_fw_utils.py

@@ -218,7 +218,23 @@ def test_create_image_uri_merged():
     image_uri = fw_utils.create_image_uri(
         "us-west-2", "mxnet-serving", "ml.c4.2xlarge", "1.4.1", "py3"
     )
-    assert image_uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-inference:1.4.1-cpu-py3"
+    assert (
+        image_uri
+        == "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.4.1-cpu-py3"
+    )
+
+    image_uri = fw_utils.create_image_uri(
+        "us-west-2",
+        "mxnet-serving",
+        "ml.c4.2xlarge",
+        "1.4.1",
+        "py3",
+        accelerator_type="ml.eia1.medium",
+    )
+    assert (
+        image_uri
+        == "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-inference-eia:1.4.1-cpu-py3"
+    )
 
 
 def test_create_image_uri_merged_py2():

tests/unit/test_mxnet.py

@@ -336,12 +336,7 @@ def test_mxnet_mms_version(
 
     model = mx.create_model()
 
-    if mxnet_version == "1.4.1":
-        expected_image_base = (
-            "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-inference:1.4.1-gpu-py2"
-        )
-    else:
-        expected_image_base = _get_full_image_uri(mxnet_version, IMAGE_REPO_SERVING_NAME, "gpu")
+    expected_image_base = _get_full_image_uri(mxnet_version, IMAGE_REPO_SERVING_NAME, "gpu")
 
     environment = {
         "Environment": {