fix: tgi image uri unit tests (#5127)

* fix: tgi image uri unit tests

* fix: black-format and flake8 failures

* fix: parse

* fix: print statement

---------

Co-authored-by: Erick Benitez-Ramos <[email protected]>

Author: evakravi

Diff for: tests/unit/sagemaker/image_uris/test_huggingface_llm.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import pytest
+from packaging.version import parse
 
 from sagemaker.huggingface import get_huggingface_llm_image_uri
 from tests.unit.sagemaker.image_uris import expected_uris, conftest
@@ -72,10 +73,31 @@ def test_huggingface_uris(load_config):
     VERSIONS = load_config["inference"]["versions"]
     device = load_config["inference"]["processors"][0]
     backend = "huggingface-neuronx" if device == "inf2" else "huggingface"
+
+    # Fail if device is not in mapping
+    if device not in HF_VERSIONS_MAPPING:
+        raise ValueError(f"Device {device} not found in HF_VERSIONS_MAPPING")
+
+    # Get highest version for the device
+    highest_version = max(HF_VERSIONS_MAPPING[device].keys(), key=lambda x: parse(x))
+
     for version in VERSIONS:
         ACCOUNTS = load_config["inference"]["versions"][version]["registries"]
         for region in ACCOUNTS.keys():
             uri = get_huggingface_llm_image_uri(backend, region=region, version=version)
+
+            # Skip only if test version is higher than highest known version.
+            # There's now automation to add new TGI releases to image_uri_config directory
+            # that doesn't involve a human raising a PR.
+            if parse(version) > parse(highest_version):
+                print(
+                    f"Skipping version check for {version} as there is "
+                    "automation that now updates the image_uri_config "
+                    "without a human raising a PR. Tests will pass for "
+                    f"versions higher than {highest_version} that are not in HF_VERSIONS_MAPPING."
+                )
+                continue
+
             expected = expected_uris.huggingface_llm_framework_uri(
                 "huggingface-pytorch-tgi-inference",
                 ACCOUNTS[region],