feature: add serverless inference image_uri support

Author: bhaoz

src/sagemaker/chainer/model.py

@@ -143,7 +143,9 @@ def __init__(
 
         self.model_server_workers = model_server_workers
 
-    def prepare_container_def(self, instance_type=None, accelerator_type=None):
+    def prepare_container_def(
+        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+    ):
         """Return a container definition with framework configuration set in model environment.
 
         Args:
@@ -159,14 +161,17 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """
         deploy_image = self.image_uri
         if not deploy_image:
-            if instance_type is None:
+            if instance_type is None and serverless_inference_config 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
+                region_name,
+                instance_type,
+                accelerator_type=accelerator_type,
+                serverless_inference_config=serverless_inference_config,
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -178,7 +183,9 @@ def prepare_container_def(self, instance_type=None, 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, accelerator_type=None):
+    def serving_image_uri(
+        self, region_name, instance_type, accelerator_type=None, serverless_inference_config=None
+    ):
         """Create a URI for the serving image.
 
         Args:
@@ -198,4 +205,5 @@ def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
             instance_type=instance_type,
             accelerator_type=accelerator_type,
             image_scope="inference",
+            serverless_inference_config=serverless_inference_config,
         )

src/sagemaker/huggingface/model.py

@@ -272,7 +272,7 @@ def deploy(
                 is not None. Otherwise, return None.
         """
 
-        if not self.image_uri and instance_type.startswith("ml.inf"):
+        if not self.image_uri and instance_type is not None and instance_type.startswith("ml.inf"):
             self.image_uri = self.serving_image_uri(
                 region_name=self.sagemaker_session.boto_session.region_name,
                 instance_type=instance_type,
@@ -365,7 +365,9 @@ def register(
             drift_check_baselines=drift_check_baselines,
         )
 
-    def prepare_container_def(self, instance_type=None, accelerator_type=None):
+    def prepare_container_def(
+        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+    ):
         """A container definition with framework configuration set in model environment variables.
 
         Args:
@@ -381,14 +383,17 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """
         deploy_image = self.image_uri
         if not deploy_image:
-            if instance_type is None:
+            if instance_type is None and serverless_inference_config 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
+                region_name,
+                instance_type,
+                accelerator_type=accelerator_type,
+                serverless_inference_config=serverless_inference_config,
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -402,7 +407,13 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
 
-    def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
+    def serving_image_uri(
+        self,
+        region_name,
+        instance_type=None,
+        accelerator_type=None,
+        serverless_inference_config=None,
+    ):
         """Create a URI for the serving image.
 
         Args:
@@ -432,4 +443,5 @@ def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
             accelerator_type=accelerator_type,
             image_scope="inference",
             base_framework_version=base_framework_version,
+            serverless_inference_config=serverless_inference_config,
         )

src/sagemaker/image_uris.py

@@ -48,6 +48,7 @@ def retrieve(
     tolerate_deprecated_model=False,
     sdk_version=None,
     inference_tool=None,
+    serverless_inference_config=None,
 ) -> str:
     """Retrieves the ECR URI for the Docker image matching the given arguments.
 
@@ -159,7 +160,9 @@ def retrieve(
     repo = version_config["repository"]
 
     processor = _processor(
-        instance_type, config.get("processors") or version_config.get("processors")
+        instance_type,
+        config.get("processors") or version_config.get("processors"),
+        serverless_inference_config,
     )
 
     # if container version is available in .json file, utilize that
@@ -202,7 +205,9 @@ def retrieve(
 
     tag = _format_tag(tag_prefix, processor, py_version, container_version, inference_tool)
 
-    if _should_auto_select_container_version(instance_type, distribution):
+    if instance_type is not None and _should_auto_select_container_version(
+        instance_type, distribution
+    ):
         container_versions = {
             "tensorflow-2.3-gpu-py37": "cu110-ubuntu18.04-v3",
             "tensorflow-2.3.1-gpu-py37": "cu110-ubuntu18.04",
@@ -327,7 +332,7 @@ def _registry_from_region(region, registry_dict):
     return registry_dict[region]
 
 
-def _processor(instance_type, available_processors):
+def _processor(instance_type, available_processors, serverless_inference_config=None):
     """Returns the processor type for the given instance type."""
     if not available_processors:
         logger.info("Ignoring unnecessary instance type: %s.", instance_type)
@@ -337,6 +342,10 @@ def _processor(instance_type, available_processors):
         logger.info("Defaulting to only supported image scope: %s.", available_processors[0])
         return available_processors[0]
 
+    if serverless_inference_config is not None:
+        logger.info("Defaulting to CPU type when using serverless inference")
+        return "cpu"
+
     if not instance_type:
         raise ValueError(
             "Empty SageMaker instance type. For options, see: "

src/sagemaker/model.py

@@ -383,7 +383,10 @@ def _init_sagemaker_session_if_does_not_exist(self, instance_type=None):
             self.sagemaker_session = session.Session()
 
     def prepare_container_def(
-        self, instance_type=None, accelerator_type=None
+        self,
+        instance_type=None,
+        accelerator_type=None,
+        serverless_inference_config=None,
     ):  # pylint: disable=unused-argument
         """Return a dict created by ``sagemaker.container_def()``.
 
@@ -498,7 +501,9 @@ def enable_network_isolation(self):
         """
         return self._enable_network_isolation
 
-    def _create_sagemaker_model(self, instance_type=None, accelerator_type=None, tags=None):
+    def _create_sagemaker_model(
+        self, instance_type=None, accelerator_type=None, tags=None, serverless_inference_config=None
+    ):
         """Create a SageMaker Model Entity
 
         Args:
@@ -515,7 +520,11 @@ def _create_sagemaker_model(self, instance_type=None, accelerator_type=None, tag
                 https://boto3.amazonaws.com/v1/documentation
                 /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
         """
-        container_def = self.prepare_container_def(instance_type, accelerator_type=accelerator_type)
+        container_def = self.prepare_container_def(
+            instance_type,
+            accelerator_type=accelerator_type,
+            serverless_inference_config=serverless_inference_config,
+        )
 
         self._ensure_base_name_if_needed(
             image_uri=container_def["Image"], script_uri=self.source_dir, model_uri=self.model_data
@@ -983,7 +992,9 @@ def deploy(
             if self._base_name is not None:
                 self._base_name = "-".join((self._base_name, compiled_model_suffix))
 
-        self._create_sagemaker_model(instance_type, accelerator_type, tags)
+        self._create_sagemaker_model(
+            instance_type, accelerator_type, tags, serverless_inference_config
+        )
 
         serverless_inference_config_dict = (
             serverless_inference_config._to_request_dict() if is_serverless else None

src/sagemaker/mxnet/model.py

@@ -220,7 +220,9 @@ def register(
             customer_metadata_properties=customer_metadata_properties,
         )
 
-    def prepare_container_def(self, instance_type=None, accelerator_type=None):
+    def prepare_container_def(
+        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+    ):
         """Return a container definition with framework configuration.
 
         Framework configuration is set in model environment variables.
@@ -238,14 +240,17 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """
         deploy_image = self.image_uri
         if not deploy_image:
-            if instance_type is None:
+            if instance_type is None and serverless_inference_config 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
+                region_name,
+                instance_type,
+                accelerator_type=accelerator_type,
+                serverless_inference_config=serverless_inference_config,
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -259,7 +264,9 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
 
-    def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
+    def serving_image_uri(
+        self, region_name, instance_type, accelerator_type=None, serverless_inference_config=None
+    ):
         """Create a URI for the serving image.
 
         Args:
@@ -282,6 +289,7 @@ def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
             instance_type=instance_type,
             accelerator_type=accelerator_type,
             image_scope="inference",
+            serverless_inference_config=serverless_inference_config,
         )
 
     def _is_mms_version(self):

src/sagemaker/pytorch/model.py

@@ -220,7 +220,9 @@ def register(
             customer_metadata_properties=customer_metadata_properties,
         )
 
-    def prepare_container_def(self, instance_type=None, accelerator_type=None):
+    def prepare_container_def(
+        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+    ):
         """A container definition with framework configuration set in model environment variables.
 
         Args:
@@ -236,14 +238,17 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """
         deploy_image = self.image_uri
         if not deploy_image:
-            if instance_type is None:
+            if instance_type is None and serverless_inference_config 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
+                region_name,
+                instance_type,
+                accelerator_type=accelerator_type,
+                serverless_inference_config=serverless_inference_config,
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -257,7 +262,9 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
 
-    def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
+    def serving_image_uri(
+        self, region_name, instance_type, accelerator_type=None, serverless_inference_config=None
+    ):
         """Create a URI for the serving image.
 
         Args:
@@ -280,6 +287,7 @@ def serving_image_uri(self, region_name, instance_type, accelerator_type=None):
             instance_type=instance_type,
             accelerator_type=accelerator_type,
             image_scope="inference",
+            serverless_inference_config=serverless_inference_config,
         )
 
     def _is_mms_version(self):

src/sagemaker/sklearn/model.py

@@ -208,7 +208,9 @@ def register(
             description,
         )
 
-    def prepare_container_def(self, instance_type=None, accelerator_type=None):
+    def prepare_container_def(
+        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+    ):
         """Container definition with framework configuration set in model environment variables.
 
         Args:
@@ -244,7 +246,7 @@ def prepare_container_def(self, instance_type=None, 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, serverless_inference_config=None):
         """Create a URI for the serving image.
 
         Args:
@@ -261,4 +263,5 @@ def serving_image_uri(self, region_name, instance_type):
             version=self.framework_version,
             py_version=self.py_version,
             instance_type=instance_type,
+            serverless_inference_config=serverless_inference_config,
         )