Pass inference_tool for generating neuron/x image_uris

Author: claytonparnell

src/sagemaker/huggingface/model.py

@@ -448,7 +448,11 @@ def register(
         )
 
     def prepare_container_def(
-        self, instance_type=None, accelerator_type=None, serverless_inference_config=None
+        self,
+        instance_type=None,
+        accelerator_type=None,
+        serverless_inference_config=None,
+        inference_tool=None,
     ):
         """A container definition with framework configuration set in model environment variables.
 
@@ -461,6 +465,8 @@ def prepare_container_def(
             serverless_inference_config (sagemaker.serverless.ServerlessInferenceConfig):
                 Specifies configuration related to serverless endpoint. Instance type is
                 not provided in serverless inference. So this is used to find image URIs.
+            inference_tool (str): the tool that will be used to aid in the inference.
+                Valid values: "neuron, neuronx, None" (default: None).
 
         Returns:
             dict[str, str]: A container definition object usable with the
@@ -479,6 +485,7 @@ def prepare_container_def(
                 instance_type,
                 accelerator_type=accelerator_type,
                 serverless_inference_config=serverless_inference_config,
+                inference_tool=inference_tool,
             )
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
@@ -500,6 +507,7 @@ def serving_image_uri(
         instance_type=None,
         accelerator_type=None,
         serverless_inference_config=None,
+        inference_tool=None,
     ):
         """Create a URI for the serving image.
 
@@ -513,6 +521,8 @@ def serving_image_uri(
             serverless_inference_config (sagemaker.serverless.ServerlessInferenceConfig):
                 Specifies configuration related to serverless endpoint. Instance type is
                 not provided in serverless inference. So this is used used to determine device type.
+            inference_tool (str): the tool that will be used to aid in the inference.
+                Valid values: "neuron, neuronx, None" (default: None).
 
         Returns:
             str: The appropriate image URI based on the given parameters.
@@ -534,4 +544,5 @@ def serving_image_uri(
             image_scope="inference",
             base_framework_version=base_framework_version,
             serverless_inference_config=serverless_inference_config,
+            inference_tool=inference_tool,
         )

src/sagemaker/image_uris.py

@@ -104,7 +104,7 @@ def retrieve(
         sdk_version (str): the version of python-sdk that will be used in the image retrieval.
             (default: None).
         inference_tool (str): the tool that will be used to aid in the inference.
-            Valid values: "neuron, None"
+            Valid values: "neuron, neuronx, None"
             (default: None).
         serverless_inference_config (sagemaker.serverless.ServerlessInferenceConfig):
             Specifies configuration related to serverless endpoint. Instance type is
@@ -158,7 +158,7 @@ def retrieve(
         _framework = framework
         if framework == HUGGING_FACE_FRAMEWORK or framework in TRAINIUM_ALLOWED_FRAMEWORKS:
             inference_tool = _get_inference_tool(inference_tool, instance_type)
-            if inference_tool == "neuron":
+            if inference_tool == "neuron" or inference_tool == "neuronx":
                 _framework = f"{framework}-{inference_tool}"
         final_image_scope = _get_final_image_scope(framework, instance_type, image_scope)
         _validate_for_suppported_frameworks_and_instance_type(framework, instance_type)

tests/unit/sagemaker/huggingface/huggingface_utils.py

@@ -50,4 +50,5 @@ def get_full_neuronx_image_uri(
         image_scope="training",
         base_framework_version=base_framework_version,
         container_version="cu110-ubuntu18.04",
+        inference_tool="neuronx",
     )

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -269,7 +269,7 @@ def test_huggingface_neuron(
         pytorch_version=huggingface_neuron_latest_inference_pytorch_version,
         py_version=huggingface_neuron_latest_inference_py_version,
     )
-    container = huggingface_model.prepare_container_def("ml.inf1.xlarge")
+    container = huggingface_model.prepare_container_def("ml.inf1.xlarge", inference_tool="neuron")
     assert container["Image"]
 
 
@@ -289,7 +289,7 @@ def test_huggingface_neuronx(
         pytorch_version=huggingface_neuronx_latest_inference_pytorch_version,
         py_version=huggingface_neuronx_latest_inference_py_version,
     )
-    container = huggingface_model.prepare_container_def("ml.inf2.xlarge")
+    container = huggingface_model.prepare_container_def("ml.inf2.xlarge", inference_tool="neuronx")
     assert container["Image"]