fix: instance type found but image uri for family

Author: evakravi

src/sagemaker/jumpstart/types.py

@@ -356,12 +356,10 @@ def get_image_uri(self, instance_type: str, region: str) -> Optional[str]:
         if None in [self.regional_aliases, self.variants]:
             return None
 
-        image_uri_alias: Optional[str] = None
-        if instance_type in self.variants:
-            image_uri_alias = (
-                self.variants[instance_type].get("regional_properties", {}).get("image_uri")
-            )
-        else:
+        image_uri_alias: Optional[str] = (
+            self.variants.get(instance_type, {}).get("regional_properties", {}).get("image_uri")
+        )
+        if image_uri_alias is None:
             instance_type_family = get_instance_type_family(instance_type)
 
             if instance_type_family in {"", None}:

tests/unit/sagemaker/jumpstart/test_types.py

@@ -175,6 +175,7 @@ def test_jumpstart_instance_variants():
                 "ml.p3.200xlarge": {"regional_properties": {"image_uri": "$gpu_image_uri_2"}},
                 "p4": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
                 "g4dn": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
+                "g9": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
                 "m2": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
                 "c2": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
                 "local": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
@@ -184,6 +185,9 @@ def test_jumpstart_instance_variants():
                 "ml.g5.12xlarge": {
                     "properties": {"environment_variables": {"TENSOR_PARALLEL_DEGREE": "4"}}
                 },
+                "ml.g9.12xlarge": {
+                    "properties": {"environment_variables": {"TENSOR_PARALLEL_DEGREE": "4"}}
+                },
             },
         }
     )
@@ -192,6 +196,12 @@ def test_jumpstart_instance_variants():
         == "763104351884.dkr.ecr.us-west-2.amazonaws.com/stud-gpu"
     )
 
+    assert (
+        variant.get_image_uri(instance_type="ml.g9.12xlarge", region="us-west-2")
+        == "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:"
+        "1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04"
+    )
+
     assert (
         variant.get_image_uri(instance_type="ml.p3.2xlarge", region="us-west-2")
         == "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:"