fix: Use serving_image_uri for Airflow (#1245)

* Add methods to framework models to get serving image URI

* Update test_airflow_config integ test

* Format with black

Author: ajaykarpur

src/sagemaker/chainer/model.py

@@ -15,6 +15,8 @@
 
 import logging
 
+from sagemaker import fw_utils
+
 import sagemaker
 from sagemaker.fw_utils import (
     create_image_uri,
@@ -154,3 +156,23 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         if self.model_server_workers:
             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):
+        """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).
+
+        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,
+        )

src/sagemaker/mxnet/model.py

@@ -16,6 +16,7 @@
 import logging
 
 from pkg_resources import parse_version
+from sagemaker import fw_utils
 
 import sagemaker
 from sagemaker.fw_utils import (
@@ -167,3 +168,23 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         return sagemaker.container_def(
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
+
+    def serving_image_uri(self, region_name, instance_type):
+        """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).
+
+        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,
+        )

src/sagemaker/pytorch/model.py

@@ -15,6 +15,7 @@
 
 import logging
 import pkg_resources
+from sagemaker import fw_utils
 
 import sagemaker
 from sagemaker.fw_utils import (
@@ -167,3 +168,23 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         return sagemaker.container_def(
             deploy_image, self.repacked_model_data or self.model_data, deploy_env
         )
+
+    def serving_image_uri(self, region_name, instance_type):
+        """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).
+
+        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,
+        )

src/sagemaker/sklearn/model.py

@@ -15,6 +15,8 @@
 
 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
@@ -151,3 +153,23 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
             self.repacked_model_data if self.enable_network_isolation() else self.model_data
         )
         return sagemaker.container_def(deploy_image, model_data_uri, deploy_env)
+
+    def serving_image_uri(self, region_name, instance_type):
+        """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).
+
+        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,
+        )

src/sagemaker/tensorflow/model.py

@@ -15,6 +15,8 @@
 
 import logging
 
+from sagemaker import fw_utils
+
 import sagemaker
 from sagemaker.fw_utils import (
     create_image_uri,
@@ -157,3 +159,23 @@ 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):
+        """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).
+
+        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,
+        )

src/sagemaker/tensorflow/serving.py

@@ -275,3 +275,17 @@ def _get_image_uri(self, instance_type, accelerator_type=None):
             self._framework_version,
             accelerator_type=accelerator_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).
+
+        Returns:
+            str: The appropriate image URI based on the given parameters.
+
+        """
+        return self._get_image_uri(instance_type=instance_type)

src/sagemaker/workflow/airflow.py

@@ -516,13 +516,7 @@ def prepare_framework_container_def(model, instance_type, s3_operations):
     deploy_image = model.image
     if not deploy_image:
         region_name = model.sagemaker_session.boto_session.region_name
-        deploy_image = fw_utils.create_image_uri(
-            region_name,
-            model.__framework_name__,
-            instance_type,
-            model.framework_version,
-            model.py_version,
-        )
+        deploy_image = model.serving_image_uri(region_name, instance_type)
 
     base_name = utils.base_name_from_image(deploy_image)
     model.name = model.name or utils.name_from_base(base_name)

src/sagemaker/xgboost/model.py

@@ -15,6 +15,8 @@
 
 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
@@ -135,3 +137,23 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         if self.model_server_workers:
             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):
+        """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).
+
+        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,
+        )