fix: account for EI and version-based ECR repo naming in serving_image_uri()

src/sagemaker/mxnet/model.py

@@ -17,8 +17,6 @@
 
 import packaging.version
 
-from sagemaker import fw_utils
-
 import sagemaker
 from sagemaker.fw_utils import (
     create_image_uri,
@@ -143,29 +141,15 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             dict[str, str]: A container definition object usable with the
             CreateModel API.
         """
-        is_mms_version = packaging.version.Version(
-            self.framework_version
-        ) >= packaging.version.Version(self._LOWEST_MMS_VERSION)
-
         deploy_image = self.image
         if not deploy_image:
             region_name = self.sagemaker_session.boto_session.region_name
-
-            framework_name = self.__framework_name__
-            if is_mms_version:
-                framework_name += "-serving"
-
-            deploy_image = create_image_uri(
-                region_name,
-                framework_name,
-                instance_type,
-                self.framework_version,
-                self.py_version,
-                accelerator_type=accelerator_type,
+            deploy_image = self.serving_image_uri(
+                region_name, instance_type, accelerator_type=accelerator_type
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
-        self._upload_code(deploy_key_prefix, is_mms_version)
+        self._upload_code(deploy_key_prefix, self._is_mms_version())
         deploy_env = dict(self.env)
         deploy_env.update(self._framework_env_vars())
 
@@ -175,22 +159,41 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
 
-    def serving_image_uri(self, region_name, instance_type):
+    def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
         """Create a URI for the serving image.
 
         Args:
             region_name (str): AWS region where the image is uploaded.
             instance_type (str): SageMaker instance type. Used to determine device type
                 (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 (default: None). For example, 'ml.eia1.medium'.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        return fw_utils.create_image_uri(
+        framework_name = self.__framework_name__
+        if self._is_mms_version():
+            framework_name = "{}-serving".format(framework_name)
+
+        return create_image_uri(
             region_name,
-            "-".join([self.__framework_name__, "serving"]),
+            framework_name,
             instance_type,
             self.framework_version,
             self.py_version,
+            accelerator_type=accelerator_type,
         )
+
+    def _is_mms_version(self):
+        """Whether the framework version corresponds to an inference image using
+        the Multi-Model Server (https://github.com/awslabs/multi-model-server).
+
+        Returns:
+            bool: If the framework version corresponds to an image using MMS.
+        """
+        lowest_mms_version = packaging.version.Version(self._LOWEST_MMS_VERSION)
+        framework_version = packaging.version.Version(self.framework_version)
+        return framework_version >= lowest_mms_version

src/sagemaker/pytorch/model.py

@@ -15,7 +15,6 @@
 
 import logging
 import packaging.version
-from sagemaker import fw_utils
 
 import sagemaker
 from sagemaker.fw_utils import (
@@ -137,34 +136,21 @@ 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. For example, 'ml.eia1.medium'.
+                model. Currently unsupported with PyTorch.
 
         Returns:
             dict[str, str]: A container definition object usable with the
             CreateModel API.
         """
-        lowest_mms_version = packaging.version.Version(self._LOWEST_MMS_VERSION)
-        framework_version = packaging.version.Version(self.framework_version)
-        is_mms_version = framework_version >= lowest_mms_version
-
         deploy_image = self.image
         if not deploy_image:
             region_name = self.sagemaker_session.boto_session.region_name
-
-            framework_name = self.__framework_name__
-            if is_mms_version:
-                framework_name += "-serving"
-
-            deploy_image = create_image_uri(
-                region_name,
-                framework_name,
-                instance_type,
-                self.framework_version,
-                self.py_version,
-                accelerator_type=accelerator_type,
+            deploy_image = self.serving_image_uri(
+                region_name, instance_type, accelerator_type=accelerator_type
             )
+
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
-        self._upload_code(deploy_key_prefix, repack=is_mms_version)
+        self._upload_code(deploy_key_prefix, repack=self._is_mms_version())
         deploy_env = dict(self.env)
         deploy_env.update(self._framework_env_vars())
 
@@ -174,22 +160,41 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
 
-    def serving_image_uri(self, region_name, instance_type):
+    def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
         """Create a URI for the serving image.
 
         Args:
             region_name (str): AWS region where the image is uploaded.
             instance_type (str): SageMaker instance type. Used to determine device type
                 (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 unsupported with PyTorch.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        return fw_utils.create_image_uri(
+        framework_name = self.__framework_name__
+        if self._is_mms_version():
+            framework_name = "{}-serving".format(framework_name)
+
+        return create_image_uri(
             region_name,
-            "-".join([self.__framework_name__, "serving"]),
+            framework_name,
             instance_type,
             self.framework_version,
             self.py_version,
+            accelerator_type=accelerator_type,
         )
+
+    def _is_mms_version(self):
+        """Whether the framework version corresponds to an inference image using
+        the Multi-Model Server (https://github.com/awslabs/multi-model-server).
+
+        Returns:
+            bool: If the framework version corresponds to an image using MMS.
+        """
+        lowest_mms_version = packaging.version.Version(self._LOWEST_MMS_VERSION)
+        framework_version = packaging.version.Version(self.framework_version)
+        return framework_version >= lowest_mms_version

src/sagemaker/tensorflow/model.py

@@ -15,8 +15,6 @@
 
 import logging
 
-from sagemaker import fw_utils
-
 import sagemaker
 from sagemaker.fw_utils import (
     create_image_uri,
@@ -146,13 +144,8 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         deploy_image = self.image
         if not deploy_image:
             region_name = self.sagemaker_session.boto_region_name
-            deploy_image = create_image_uri(
-                region_name,
-                self.__framework_name__,
-                instance_type,
-                self.framework_version,
-                self.py_version,
-                accelerator_type=accelerator_type,
+            deploy_image = self.serving_image_uri(
+                region_name, instance_type, accelerator_type=accelerator_type
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -165,22 +158,26 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
 
         return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
 
-    def serving_image_uri(self, region_name, instance_type):
+    def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
         """Create a URI for the serving image.
 
         Args:
             region_name (str): AWS region where the image is uploaded.
             instance_type (str): SageMaker instance type. Used to determine device type
                 (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 (default: None). For example, 'ml.eia1.medium'.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        return fw_utils.create_image_uri(
+        return create_image_uri(
             region_name,
-            "-".join([self.__framework_name__, "serving"]),
+            self.__framework_name__,
             instance_type,
             self.framework_version,
             self.py_version,
+            accelerator_type=accelerator_type,
         )

src/sagemaker/tensorflow/serving.py

@@ -276,16 +276,21 @@ def _get_image_uri(self, instance_type, accelerator_type=None):
             accelerator_type=accelerator_type,
         )
 
-    def serving_image_uri(self, region_name, instance_type):  # pylint: disable=unused-argument
+    def serving_image_uri(
+        self, region_name, instance_type, accelerator_type=None
+    ):  # pylint: disable=unused-argument
         """Create a URI for the serving image.
 
         Args:
             region_name (str): AWS region where the image is uploaded.
             instance_type (str): SageMaker instance type. Used to determine device type
                 (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 (default: None). For example, 'ml.eia1.medium'.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        return self._get_image_uri(instance_type=instance_type)
+        return self._get_image_uri(instance_type=instance_type, accelerator_type=accelerator_type)

tests/unit/test_airflow.py

@@ -1030,7 +1030,7 @@ def test_model_config_from_framework_estimator(sagemaker_session):
         entry_point="{{ entry_point }}",
         source_dir="{{ source_dir }}",
         py_version="py3",
-        framework_version="1.3.0",
+        framework_version="1.6.0",
         role="{{ role }}",
         train_instance_count=1,
         train_instance_type="ml.m4.xlarge",
@@ -1051,9 +1051,9 @@ def test_model_config_from_framework_estimator(sagemaker_session):
         task_type="training",
     )
     expected_config = {
-        "ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+        "ModelName": "mxnet-inference-%s" % TIME_STAMP,
         "PrimaryContainer": {
-            "Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.3.0-cpu-py3",
+            "Image": "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-inference:1.6.0-cpu-py3",
             "Environment": {
                 "SAGEMAKER_PROGRAM": "{{ entry_point }}",
                 "SAGEMAKER_SUBMIT_DIRECTORY": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']"
@@ -1184,7 +1184,7 @@ def test_transform_config_from_framework_estimator(sagemaker_session):
         entry_point="{{ entry_point }}",
         source_dir="{{ source_dir }}",
         py_version="py3",
-        framework_version="1.3.0",
+        framework_version="1.6.0",
         role="{{ role }}",
         train_instance_count=1,
         train_instance_type="ml.m4.xlarge",
@@ -1209,9 +1209,9 @@ def test_transform_config_from_framework_estimator(sagemaker_session):
     )
     expected_config = {
         "Model": {
-            "ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+            "ModelName": "mxnet-inference-%s" % TIME_STAMP,
             "PrimaryContainer": {
-                "Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.3.0-gpu-py3",
+                "Image": "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-inference:1.6.0-gpu-py3",
                 "Environment": {
                     "SAGEMAKER_PROGRAM": "{{ entry_point }}",
                     "SAGEMAKER_SUBMIT_DIRECTORY": "s3://output/{{ ti.xcom_pull(task_ids='task_id')"
@@ -1228,7 +1228,7 @@ def test_transform_config_from_framework_estimator(sagemaker_session):
         },
         "Transform": {
             "TransformJobName": "{{ base_job_name }}-%s" % TIME_STAMP,
-            "ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+            "ModelName": "mxnet-inference-%s" % TIME_STAMP,
             "TransformInput": {
                 "DataSource": {
                     "S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": "{{ transform_data }}"}
@@ -1425,7 +1425,7 @@ def test_deploy_config_from_framework_estimator(sagemaker_session):
         entry_point="{{ entry_point }}",
         source_dir="{{ source_dir }}",
         py_version="py3",
-        framework_version="1.3.0",
+        framework_version="1.6.0",
         role="{{ role }}",
         train_instance_count=1,
         train_instance_type="ml.m4.xlarge",
@@ -1449,9 +1449,9 @@ def test_deploy_config_from_framework_estimator(sagemaker_session):
     )
     expected_config = {
         "Model": {
-            "ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+            "ModelName": "mxnet-inference-%s" % TIME_STAMP,
             "PrimaryContainer": {
-                "Image": "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-serving:1.3.0-cpu-py3",
+                "Image": "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-inference:1.6.0-cpu-py3",
                 "Environment": {
                     "SAGEMAKER_PROGRAM": "{{ entry_point }}",
                     "SAGEMAKER_SUBMIT_DIRECTORY": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']"
@@ -1466,20 +1466,20 @@ def test_deploy_config_from_framework_estimator(sagemaker_session):
             "ExecutionRoleArn": "{{ role }}",
         },
         "EndpointConfig": {
-            "EndpointConfigName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+            "EndpointConfigName": "mxnet-inference-%s" % TIME_STAMP,
             "ProductionVariants": [
                 {
                     "InstanceType": "ml.c4.large",
                     "InitialInstanceCount": "{{ instance_count}}",
-                    "ModelName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+                    "ModelName": "mxnet-inference-%s" % TIME_STAMP,
                     "VariantName": "AllTraffic",
                     "InitialVariantWeight": 1,
                 }
             ],
         },
         "Endpoint": {
             "EndpointName": "mxnet-endpoint",
-            "EndpointConfigName": "sagemaker-mxnet-serving-%s" % TIME_STAMP,
+            "EndpointConfigName": "mxnet-inference-%s" % TIME_STAMP,
         },
     }