change: make instance_type optional for prepare_container_def

src/sagemaker/chainer/model.py

@@ -15,8 +15,6 @@
 
 import logging
 
-from sagemaker import fw_utils
-
 import sagemaker
 from sagemaker.fw_utils import (
     create_image_uri,
@@ -126,7 +124,7 @@ def __init__(
         self.framework_version = framework_version or defaults.CHAINER_VERSION
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.
 
@@ -143,14 +141,14 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         """
         deploy_image = self.image
         if not deploy_image:
+            if instance_type is None:
+                raise ValueError(
+                    "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+                )
+
             region_name = self.sagemaker_session.boto_session.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)
@@ -162,7 +160,7 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
         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:
@@ -174,10 +172,11 @@ def serving_image_uri(self, region_name, instance_type):
             str: The appropriate image URI based on the given parameters.
 
         """
-        return fw_utils.create_image_uri(
+        return create_image_uri(
             region_name,
             self.__framework_name__,
             instance_type,
             self.framework_version,
             self.py_version,
+            accelerator_type=accelerator_type,
         )

src/sagemaker/model.py

@@ -138,7 +138,7 @@ def _init_sagemaker_session_if_does_not_exist(self, instance_type):
             self.sagemaker_session = session.Session()
 
     def prepare_container_def(
-        self, instance_type, accelerator_type=None
+        self, instance_type=None, accelerator_type=None
     ):  # pylint: disable=unused-argument
         """Return a dict created by ``sagemaker.container_def()`` for deploying
         this model to a specified instance type.
@@ -166,7 +166,7 @@ def enable_network_isolation(self):
         """
         return self._enable_network_isolation
 
-    def _create_sagemaker_model(self, instance_type, accelerator_type=None, tags=None):
+    def _create_sagemaker_model(self, instance_type=None, accelerator_type=None, tags=None):
         """Create a SageMaker Model Entity
 
         Args:
@@ -807,9 +807,7 @@ def __init__(
         self.uploaded_code = None
         self.repacked_model_data = None
 
-    def prepare_container_def(
-        self, instance_type, accelerator_type=None
-    ):  # pylint disable=unused-argument
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.
 

src/sagemaker/multidatamodel.py

@@ -111,7 +111,7 @@ def __init__(
                 **kwargs
             )
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition set with MultiModel mode,
         model data and other parameters from the model (if available).
 

src/sagemaker/mxnet/model.py

@@ -126,7 +126,7 @@ def __init__(
         self.framework_version = framework_version or defaults.MXNET_VERSION
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.
 
@@ -143,6 +143,11 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         """
         deploy_image = self.image
         if not deploy_image:
+            if instance_type is None:
+                raise ValueError(
+                    "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+                )
+
             region_name = self.sagemaker_session.boto_session.region_name
             deploy_image = self.serving_image_uri(
                 region_name, instance_type, accelerator_type=accelerator_type

src/sagemaker/pytorch/model.py

@@ -127,7 +127,7 @@ def __init__(
         self.framework_version = framework_version or defaults.PYTORCH_VERSION
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.
 
@@ -144,6 +144,11 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         """
         deploy_image = self.image
         if not deploy_image:
+            if instance_type is None:
+                raise ValueError(
+                    "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+                )
+
             region_name = self.sagemaker_session.boto_session.region_name
             deploy_image = self.serving_image_uri(
                 region_name, instance_type, accelerator_type=accelerator_type

src/sagemaker/sklearn/model.py

@@ -15,8 +15,6 @@
 
 import logging
 
-from sagemaker import fw_utils
-
 import sagemaker
 from sagemaker.fw_utils import model_code_key_prefix, python_deprecation_warning
 from sagemaker.fw_registry import default_framework_uri
@@ -118,16 +116,16 @@ def __init__(
         self.framework_version = framework_version
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.
 
         Args:
             instance_type (str): The EC2 instance type to deploy this Model to.
-                For example, 'ml.p2.xlarge'.
+                This parameter is unused because Scikit-learn supports only CPU.
             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'. Note: accelerator types
+                model. This parameter is unused because accelerator types
                 are not supported by SKLearnModel.
 
         Returns:
@@ -139,9 +137,8 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
 
         deploy_image = self.image
         if not deploy_image:
-            image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
-            deploy_image = default_framework_uri(
-                self.__framework_name__, self.sagemaker_session.boto_region_name, image_tag
+            deploy_image = self.serving_image_uri(
+                self.sagemaker_session.boto_region_name, instance_type
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -156,22 +153,17 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         )
         return sagemaker.container_def(deploy_image, model_data_uri, deploy_env)
 
-    def serving_image_uri(self, region_name, instance_type):
+    def serving_image_uri(self, region_name, instance_type):  # 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).
+            instance_type (str): SageMaker instance type. This parameter is unused because
+                Scikit-learn supports only CPU.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        return fw_utils.create_image_uri(
-            region_name,
-            self.__framework_name__,
-            instance_type,
-            self.framework_version,
-            self.py_version,
-        )
+        image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
+        return default_framework_uri(self.__framework_name__, region_name, image_tag)

src/sagemaker/tensorflow/model.py

@@ -124,7 +124,7 @@ def __init__(
         self.framework_version = framework_version or defaults.TF_VERSION
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.
 
@@ -143,6 +143,11 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         """
         deploy_image = self.image
         if not deploy_image:
+            if instance_type is None:
+                raise ValueError(
+                    "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+                )
+
             region_name = self.sagemaker_session.boto_region_name
             deploy_image = self.serving_image_uri(
                 region_name, instance_type, accelerator_type=accelerator_type

src/sagemaker/tensorflow/serving.py

@@ -215,12 +215,17 @@ 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
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """
         Args:
             instance_type:
             accelerator_type:
         """
+        if self.image is None and instance_type is None:
+            raise ValueError(
+                "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+            )
+
         image = self._get_image_uri(instance_type, accelerator_type)
         env = self._get_container_env()
 

src/sagemaker/xgboost/model.py

@@ -15,8 +15,6 @@
 
 import logging
 
-from sagemaker import fw_utils
-
 import sagemaker
 from sagemaker.fw_utils import model_code_key_prefix
 from sagemaker.fw_registry import default_framework_uri
@@ -107,26 +105,24 @@ def __init__(
         self.framework_version = framework_version
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type, accelerator_type=None):
+    def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration
         set in model environment variables.
 
         Args:
-            instance_type (str): The EC2 instance type to deploy this Model to. For example,
-                'ml.m5.xlarge'.
+            instance_type (str): The EC2 instance type to deploy this Model to.
+                This parameter is unused because XGBoost supports only CPU.
             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'.
-            Note: accelerator types are not supported by XGBoostModel.
+            instance for loading and making inferences to the model. This parameter is
+                unused because accelerator types are not supported by XGBoostModel.
 
         Returns:
             dict[str, str]: A container definition object usable with the CreateModel API.
         """
         deploy_image = self.image
         if not deploy_image:
-            image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
-            deploy_image = default_framework_uri(
-                self.__framework_name__, self.sagemaker_session.boto_region_name, image_tag
+            deploy_image = self.serving_image_uri(
+                self.sagemaker_session.boto_region_name, instance_type
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -138,22 +134,17 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
         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):  # 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).
+            instance_type (str): SageMaker instance type. This parameter is unused because
+                XGBoost supports only CPU.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        return fw_utils.create_image_uri(
-            region_name,
-            self.__framework_name__,
-            instance_type,
-            self.framework_version,
-            self.py_version,
-        )
+        image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
+        return default_framework_uri(self.__framework_name__, region_name, image_tag)

tests/unit/sagemaker/model/test_model.py

@@ -37,9 +37,12 @@ def test_prepare_container_def():
     env = {"FOO": "BAR"}
     model = Model(MODEL_DATA, MODEL_IMAGE, env=env)
 
+    expected = {"Image": MODEL_IMAGE, "Environment": env, "ModelDataUrl": MODEL_DATA}
+
     container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
+    assert expected == container_def
 
-    expected = {"Image": MODEL_IMAGE, "Environment": env, "ModelDataUrl": MODEL_DATA}
+    container_def = model.prepare_container_def()
     assert expected == container_def
 
 
@@ -60,16 +63,25 @@ def test_create_sagemaker_model(name_from_image, prepare_container_def, sagemake
     prepare_container_def.return_value = container_def
 
     model = Model(MODEL_DATA, MODEL_IMAGE, sagemaker_session=sagemaker_session)
-    model._create_sagemaker_model(INSTANCE_TYPE)
+    model._create_sagemaker_model()
 
-    prepare_container_def.assert_called_with(INSTANCE_TYPE, accelerator_type=None)
+    prepare_container_def.assert_called_with(None, accelerator_type=None)
     name_from_image.assert_called_with(MODEL_IMAGE)
 
     sagemaker_session.create_model.assert_called_with(
         MODEL_NAME, None, container_def, vpc_config=None, enable_network_isolation=False, tags=None
     )
 
 
+@patch("sagemaker.utils.name_from_image", Mock())
+@patch("sagemaker.model.Model.prepare_container_def")
+def test_create_sagemaker_model_instance_type(prepare_container_def, sagemaker_session):
+    model = Model(MODEL_DATA, MODEL_IMAGE, sagemaker_session=sagemaker_session)
+    model._create_sagemaker_model(INSTANCE_TYPE)
+
+    prepare_container_def.assert_called_with(INSTANCE_TYPE, accelerator_type=None)
+
+
 @patch("sagemaker.utils.name_from_image", Mock())
 @patch("sagemaker.model.Model.prepare_container_def")
 def test_create_sagemaker_model_accelerator_type(prepare_container_def, sagemaker_session):

tests/unit/test_chainer.py

@@ -434,6 +434,16 @@ def test_model_prepare_container_def_accelerator_error(sagemaker_session):
         model.prepare_container_def(INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
 
 
+def test_model_prepare_container_def_no_instance_type_or_image():
+    model = ChainerModel(MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH)
+
+    with pytest.raises(ValueError) as e:
+        model.prepare_container_def()
+
+    expected_msg = "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+    assert expected_msg in str(e)
+
+
 def test_train_image_default(sagemaker_session):
     chainer = Chainer(
         entry_point=SCRIPT_PATH,

tests/unit/test_mxnet.py

@@ -491,6 +491,16 @@ def test_model_image_accelerator_mms_version(sagemaker_session):
     )
 
 
+def test_model_prepare_container_def_no_instance_type_or_image():
+    model = MXNetModel(MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH)
+
+    with pytest.raises(ValueError) as e:
+        model.prepare_container_def()
+
+    expected_msg = "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
+    assert expected_msg in str(e)
+
+
 def test_train_image_default(sagemaker_session):
     mx = MXNet(
         entry_point=SCRIPT_PATH,