fix: update image and hardware validation with inf and graviton

src/sagemaker/serve/utils/types.py

@@ -40,3 +40,6 @@ def __str__(self) -> str:
 
     CPU = 1
     GPU = 2
+    INFERENTIA_1 = 3
+    INFERENTIA_2 = 4
+    GRAVITON = 5

src/sagemaker/serve/validations/check_image_and_hardware_type.py

@@ -17,6 +17,21 @@
 }
 
 
+INF1_INSTANCE_FAMILIES = {"ml.inf1"}
+INF2_INSTANCE_FAMILIES = {"ml.inf2"}
+
+GRAVITON_INSTANCE_FAMILIES = {
+    "ml.c7g",
+    "ml.m6g",
+    "ml.m6gd",
+    "ml.c6g",
+    "ml.c6gd",
+    "ml.c6gn",
+    "ml.r6g",
+    "ml.r6gd",
+}
+
+
 def validate_image_uri_and_hardware(image_uri: str, instance_type: str, model_server: ModelServer):
     """Placeholder docstring"""
     if "xgboost" in image_uri:
@@ -57,6 +72,12 @@ def detect_hardware_type_of_instance(instance_type: str) -> HardwareType:
     instance_family = instance_type.rsplit(".", 1)[0]
     if instance_family in GPU_INSTANCE_FAMILIES:
         return HardwareType.GPU
+    if instance_family in INF1_INSTANCE_FAMILIES:
+        return HardwareType.INFERENTIA_1
+    if instance_family in INF2_INSTANCE_FAMILIES:
+        return HardwareType.INFERENTIA_2
+    if instance_family in GRAVITON_INSTANCE_FAMILIES:
+        return HardwareType.GRAVITON
     return HardwareType.CPU
 
 
@@ -67,4 +88,13 @@ def detect_triton_image_hardware_type(image_uri: str) -> HardwareType:
 
 def detect_torchserve_image_hardware_type(image_uri: str) -> HardwareType:
     """Placeholder docstring"""
-    return HardwareType.CPU if "cpu" in image_uri else HardwareType.GPU
+    if "neuronx" in image_uri:
+        return HardwareType.INFERENTIA_2
+    if "neuron" in image_uri:
+        return HardwareType.INFERENTIA_1
+    if "graviton" in image_uri:
+        return HardwareType.GRAVITON
+    if "cpu" in image_uri:
+        return HardwareType.CPU
+
+    return HardwareType.GPU

tests/unit/sagemaker/serve/validations/test_check_image_and_hardware_type.py

@@ -31,9 +31,25 @@
     "301217895009.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tritonserver:23.08-py3-cpu"
 )
 GPU_IMAGE_TRITON = "301217895009.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tritonserver:23.08-py3"
+GRAVITON_IMAGE_TORCHSERVE = (
+    "763104351884.dkr.ecr.us-east-1.amazonaws.com/"
+    "pytorch-inference-graviton:2.1.0-cpu-py310-ubuntu20.04-sagemaker"
+)
+INF1_IMAGE_TORCHSERVE = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com"
+    "/pytorch-inference-neuron:1.13.1-neuron-py310-sdk2.15.0-ubuntu20.04"
+)
+
+INF2_IMAGE_TORCHSERVE = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com"
+    "/pytorch-inference-neuronx:1.13.1-neuronx-py310-sdk2.15.0-ubuntu20.04"
+)
 
 CPU_INSTANCE = "ml.c5.xlarge"
 GPU_INSTANCE = "ml.g4dn.xlarge"
+INF1_INSTANCE = "ml.inf1.xlarge"
+INF2_INSTANCE = "ml.inf2.xlarge"
+GRAVITON_INSTANCE = "ml.c7g.xlarge"
 
 
 class TestValidateImageAndHardware(unittest.TestCase):
@@ -116,3 +132,53 @@ def test_triton_gpu_image_with_cpu_instance(self):
             )
 
             mock_logger.assert_called_once()
+
+    def test_torchserve_inf1_image_with_inf1_instance(self):
+
+        with patch("logging.Logger.warning") as mock_logger:
+            validate_image_uri_and_hardware(
+                image_uri=INF1_IMAGE_TORCHSERVE,
+                instance_type=INF1_INSTANCE,
+                model_server=ModelServer.TORCHSERVE,
+            )
+            mock_logger.assert_not_called()
+
+    def test_torchserve_inf2_image_with_inf2_instance(self):
+
+        with patch("logging.Logger.warning") as mock_logger:
+            validate_image_uri_and_hardware(
+                image_uri=INF2_IMAGE_TORCHSERVE,
+                instance_type=INF2_INSTANCE,
+                model_server=ModelServer.TORCHSERVE,
+            )
+            mock_logger.assert_not_called()
+
+    def test_torchserve_graviton_image_with_graviton_instance(self):
+
+        with patch("logging.Logger.warning") as mock_logger:
+            validate_image_uri_and_hardware(
+                image_uri=GRAVITON_IMAGE_TORCHSERVE,
+                instance_type=GRAVITON_INSTANCE,
+                model_server=ModelServer.TORCHSERVE,
+            )
+            mock_logger.assert_not_called()
+
+    def test_torchserve_inf1_image_with_cpu_instance(self):
+
+        with patch("logging.Logger.warning") as mock_logger:
+            validate_image_uri_and_hardware(
+                image_uri=INF1_IMAGE_TORCHSERVE,
+                instance_type=CPU_INSTANCE,
+                model_server=ModelServer.TORCHSERVE,
+            )
+            mock_logger.assert_called_once()
+
+    def test_torchserve_graviton_image_with_cpu_instance(self):
+
+        with patch("logging.Logger.warning") as mock_logger:
+            validate_image_uri_and_hardware(
+                image_uri=GRAVITON_IMAGE_TORCHSERVE,
+                instance_type=CPU_INSTANCE,
+                model_server=ModelServer.TORCHSERVE,
+            )
+            mock_logger.assert_called_once()