fix: Neuronx image retrieval missing sdk information (#3948)

Co-authored-by: Zuoyuan Huang <[email protected]>

Author: zuoyuanh

src/sagemaker/image_uris.py

@@ -210,7 +210,10 @@ def retrieve(
             "huggingface-tensorflow-trcomp-training",
         ]:
             _version = version
-        if repo in ["huggingface-pytorch-inference-neuron"]:
+        if repo in [
+            "huggingface-pytorch-inference-neuron",
+            "huggingface-pytorch-inference-neuronx",
+        ]:
             if not sdk_version:
                 sdk_version = _get_latest_versions(version_config["sdk_versions"])
             container_version = sdk_version + "-" + container_version

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -291,6 +291,7 @@ def test_huggingface_neuronx(
     )
     container = huggingface_model.prepare_container_def("ml.inf2.xlarge", inference_tool="neuronx")
     assert container["Image"]
+    assert "sdk" in container["Image"] and "py" in container["Image"]
 
 
 def test_attach(