change: use image_uris.retrieve for Neo and Inferentia images

src/sagemaker/image_uri_config/inferentia-mxnet.json

@@ -0,0 +1,14 @@
+{
+    "processors": ["inf"],
+    "scope": ["inference"],
+    "versions": {
+        "1.5.1": {
+            "py_versions": ["py3"],
+            "registries": {
+                "us-east-1": "785573368785",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-neo-mxnet"
+        }
+    }
+}

src/sagemaker/image_uri_config/inferentia-tensorflow.json

@@ -0,0 +1,14 @@
+{
+    "processors": ["inf"],
+    "scope": ["inference"],
+    "versions": {
+        "1.15.0": {
+            "py_versions": ["py3"],
+            "registries": {
+                "us-east-1": "785573368785",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-neo-tensorflow"
+        }
+    }
+}

src/sagemaker/image_uri_config/neo-mxnet.json

@@ -0,0 +1,46 @@
+{
+    "processors": ["cpu", "gpu"],
+    "scope": ["inference"],
+    "version_aliases": {
+        "0.12.1": "1.5",
+        "1.0.0": "1.5",
+        "1.1.0": "1.5",
+        "1.2": "1.5",
+        "1.2.0": "1.5",
+        "1.2.1": "1.5",
+        "1.3": "1.5",
+        "1.3.0": "1.5",
+        "1.4": "1.5",
+        "1.4.0": "1.5",
+        "1.4.1": "1.5"
+    },
+    "versions": {
+        "1.5": {
+            "py_versions": ["py3"],
+            "registries": {
+                "ap-east-1": "110948597952",
+                "ap-northeast-1": "941853720454",
+                "ap-northeast-2": "151534178276",
+                "ap-south-1": "763008648453",
+                "ap-southeast-1": "324986816169",
+                "ap-southeast-2": "355873309152",
+                "ca-central-1": "464438896020",
+                "cn-north-1": "472730292857",
+                "cn-northwest-1": "474822919863",
+                "eu-central-1": "746233611703",
+                "eu-north-1": "601324751636",
+                "eu-west-1": "802834080501",
+                "eu-west-2": "205493899709",
+                "eu-west-3": "254080097072",
+                "me-south-1": "836785723513",
+                "sa-east-1": "756306329178",
+                "us-east-1": "785573368785",
+                "us-east-2": "007439368137",
+                "us-gov-west-1": "263933020539",
+                "us-west-1": "710691900526",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-neo-mxnet"
+        }
+    }
+}

src/sagemaker/image_uri_config/neo-pytorch.json

@@ -0,0 +1,40 @@
+{
+    "processors": ["cpu", "gpu"],
+    "scope": ["inference"],
+    "version_aliases": {
+        "0.4.0": "1.4.0",
+        "1.0.0": "1.4.0",
+        "1.1.0": "1.4.0",
+        "1.2.0": "1.4.0",
+        "1.3.0": "1.4.0"
+    },
+    "versions": {
+        "1.4.0": {
+            "py_versions": ["py3"],
+            "registries": {
+                "ap-east-1": "110948597952",
+                "ap-northeast-1": "941853720454",
+                "ap-northeast-2": "151534178276",
+                "ap-south-1": "763008648453",
+                "ap-southeast-1": "324986816169",
+                "ap-southeast-2": "355873309152",
+                "ca-central-1": "464438896020",
+                "cn-north-1": "472730292857",
+                "cn-northwest-1": "474822919863",
+                "eu-central-1": "746233611703",
+                "eu-north-1": "601324751636",
+                "eu-west-1": "802834080501",
+                "eu-west-2": "205493899709",
+                "eu-west-3": "254080097072",
+                "me-south-1": "836785723513",
+                "sa-east-1": "756306329178",
+                "us-east-1": "785573368785",
+                "us-east-2": "007439368137",
+                "us-gov-west-1": "263933020539",
+                "us-west-1": "710691900526",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-neo-pytorch"
+        }
+    }
+}

src/sagemaker/image_uri_config/neo-tensorflow.json

@@ -0,0 +1,46 @@
+{
+    "processors": ["cpu", "gpu"],
+    "scope": ["inference"],
+    "version_aliases": {
+        "1.4.1": "1.15.0",
+        "1.5.0": "1.15.0",
+        "1.6.0": "1.15.0",
+        "1.7.0": "1.15.0",
+        "1.8.0": "1.15.0",
+        "1.9.0": "1.15.0",
+        "1.10.0": "1.15.0",
+        "1.11.0": "1.15.0",
+        "1.12.0": "1.15.0",
+        "1.13.0": "1.15.0",
+        "1.14.0": "1.15.0"
+    },
+    "versions": {
+        "1.15.0": {
+            "py_versions": ["py3"],
+            "registries": {
+                "ap-east-1": "110948597952",
+                "ap-northeast-1": "941853720454",
+                "ap-northeast-2": "151534178276",
+                "ap-south-1": "763008648453",
+                "ap-southeast-1": "324986816169",
+                "ap-southeast-2": "355873309152",
+                "ca-central-1": "464438896020",
+                "cn-north-1": "472730292857",
+                "cn-northwest-1": "474822919863",
+                "eu-central-1": "746233611703",
+                "eu-north-1": "601324751636",
+                "eu-west-1": "802834080501",
+                "eu-west-2": "205493899709",
+                "eu-west-3": "254080097072",
+                "me-south-1": "836785723513",
+                "sa-east-1": "756306329178",
+                "us-east-1": "785573368785",
+                "us-east-2": "007439368137",
+                "us-gov-west-1": "263933020539",
+                "us-west-1": "710691900526",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-neo-tensorflow"
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import os
+import re
 
 from sagemaker import utils
 
@@ -166,14 +167,27 @@ def _processor(instance_type, available_processors):
 
     if instance_type.startswith("local"):
         processor = "cpu" if instance_type == "local" else "gpu"
-    elif not instance_type.startswith("ml."):
-        raise ValueError(
-            "Invalid SageMaker instance type: {}. For options, see: "
-            "https://aws.amazon.com/sagemaker/pricing/instance-types".format(instance_type)
-        )
     else:
-        family = instance_type.split(".")[1]
-        processor = "gpu" if family[0] in ("g", "p") else "cpu"
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match:
+            family = match[1]
+
+            # For some frameworks, we have optimized images for specific families, e.g c5 or p3.
+            # In those cases, we use the family name in the image tag. In other cases, we use
+            # 'cpu' or 'gpu'.
+            if family in available_processors:
+                processor = family
+            elif family.startswith("inf"):
+                processor = "inf"
+            elif family[0] in ("g", "p"):
+                processor = "gpu"
+            else:
+                processor = "cpu"
+        else:
+            raise ValueError(
+                "Invalid SageMaker instance type: {}. For options, see: "
+                "https://aws.amazon.com/sagemaker/pricing/instance-types".format(instance_type)
+            )
 
     _validate_arg(processor, available_processors, "processor")
     return processor

src/sagemaker/model.py

@@ -18,7 +18,7 @@
 import os
 
 import sagemaker
-from sagemaker import fw_utils, local, session, utils, git_utils
+from sagemaker import fw_utils, image_uris, local, session, utils, git_utils
 from sagemaker.fw_utils import UploadedCode
 from sagemaker.transformer import Transformer
 
@@ -28,32 +28,6 @@
     ["mxnet", "tensorflow", "keras", "pytorch", "onnx", "xgboost", "tflite"]
 )
 
-NEO_IMAGE_ACCOUNT = {
-    "us-west-1": "710691900526",
-    "us-west-2": "301217895009",
-    "us-east-1": "785573368785",
-    "us-east-2": "007439368137",
-    "eu-west-1": "802834080501",
-    "eu-west-2": "205493899709",
-    "eu-west-3": "254080097072",
-    "eu-central-1": "746233611703",
-    "eu-north-1": "601324751636",
-    "ap-northeast-1": "941853720454",
-    "ap-northeast-2": "151534178276",
-    "ap-east-1": "110948597952",
-    "ap-southeast-1": "324986816169",
-    "ap-southeast-2": "355873309152",
-    "ap-south-1": "763008648453",
-    "sa-east-1": "756306329178",
-    "ca-central-1": "464438896020",
-    "me-south-1": "836785723513",
-    "cn-north-1": "472730292857",
-    "cn-northwest-1": "474822919863",
-    "us-gov-west-1": "263933020539",
-}
-
-INFERENTIA_INSTANCE_PREFIX = "ml_inf"
-
 
 class Model(object):
     """A SageMaker ``Model`` that can be deployed to an ``Endpoint``."""
@@ -243,7 +217,7 @@ def _compilation_job_config(
             "DataInputConfig": input_shape
             if not isinstance(input_shape, dict)
             else json.dumps(input_shape),
-            "Framework": framework,
+            "Framework": framework.upper(),
         }
         role = self.sagemaker_session.expand_role(role)
         output_model_config = {
@@ -260,64 +234,23 @@ def _compilation_job_config(
             "job_name": job_name,
         }
 
-    def check_neo_region(self, region):
-        """Check if this ``Model`` in the available region where neo support.
-
-        Args:
-            region (str): Specifies the region where want to execute compilation
-
-        Returns:
-            bool: boolean value whether if neo is available in the specified
-            region
-        """
-        if region in NEO_IMAGE_ACCOUNT:
-            return True
-        return False
+    def _compilation_image_uri(self, region, target_instance_type, framework, framework_version):
+        """Retrieve the Neo or Inferentia image URI.
 
-    def _neo_image_account(self, region):
-        """
-        Args:
-            region:
-        """
-        if region not in NEO_IMAGE_ACCOUNT:
-            raise ValueError(
-                "Neo is not currently supported in {}, "
-                "valid regions: {}".format(region, NEO_IMAGE_ACCOUNT.keys())
-            )
-        return NEO_IMAGE_ACCOUNT[region]
-
-    def _neo_image_uri(self, region, target_instance_type, framework, framework_version):
-        """
         Args:
-            region:
-            target_instance_type:
-            framework:
-            framework_version:
-        """
-        return fw_utils.create_image_uri(
-            region,
-            "neo-" + framework.lower(),
-            target_instance_type.replace("_", "."),
-            framework_version,
-            py_version="py3",
-            account=self._neo_image_account(region),
-        )
-
-    def _inferentia_image_uri(self, region, target_instance_type, framework, framework_version):
+            region (str): The AWS region.
+            target_instance_type (str): Identifies the device on which you want to run
+                your model after compilation, for example: ml_c5. For valid values, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
+            framework (str): The framework name.
+            framework_version (str): The framework version.
         """
-                Args:
-                    region:
-                    target_instance_type:
-                    framework:
-                    framework_version:
-                """
-        return fw_utils.create_image_uri(
+        framework_prefix = "inferentia-" if target_instance_type.startswith("ml_inf") else "neo-"
+        return image_uris.retrieve(
+            "{}{}".format(framework_prefix, framework),
             region,
-            "neo-" + framework.lower(),
-            target_instance_type.replace("_", "."),
-            framework_version,
-            py_version="py3",
-            account=self._neo_image_account(region),
+            instance_type=target_instance_type,
+            version=framework_version,
         )
 
     def compile(
@@ -361,7 +294,7 @@ def compile(
             sagemaker.model.Model: A SageMaker ``Model`` object. See
             :func:`~sagemaker.model.Model` for full details.
         """
-        framework = self._framework() or framework
+        framework = framework or self._framework()
         if framework is None:
             raise ValueError(
                 "You must specify framework, allowed values {}".format(NEO_ALLOWED_FRAMEWORKS)
@@ -375,8 +308,7 @@ def compile(
         if self.model_data is None:
             raise ValueError("You must provide an S3 path to the compressed model artifacts.")
 
-        framework = framework.upper()
-        framework_version = self._get_framework_version() or framework_version
+        framework_version = framework_version or self._get_framework_version()
 
         self._init_sagemaker_session_if_does_not_exist(target_instance_family)
         config = self._compilation_job_config(
@@ -392,16 +324,9 @@ def compile(
         self.sagemaker_session.compile_model(**config)
         job_status = self.sagemaker_session.wait_for_compilation_job(job_name)
         self.model_data = job_status["ModelArtifacts"]["S3ModelArtifacts"]
+
         if target_instance_family.startswith("ml_"):
-            self.image_uri = self._neo_image_uri(
-                self.sagemaker_session.boto_region_name,
-                target_instance_family,
-                framework,
-                framework_version,
-            )
-            self._is_compiled_model = True
-        elif target_instance_family.startswith(INFERENTIA_INSTANCE_PREFIX):
-            self.image_uri = self._inferentia_image_uri(
+            self.image_uri = self._compilation_image_uri(
                 self.sagemaker_session.boto_region_name,
                 target_instance_family,
                 framework,
@@ -414,6 +339,7 @@ def compile(
                 "please deploy the model manually.",
                 target_instance_family,
             )
+
         return self
 
     def deploy(