Fix: remove the special condition and fix the unit test (#1601)

* Fix: move the functionality from latest_container_image to retrieve

* address some comments from Gokul and add unit test

* remove extra functions and rewrite the test

* fix unit test

* fix for other unit test

* unit test fix

* fix unit test: add one more condition

* more unit tests fix

* remove redundant files

* remove the special condition and fix the unit test

---------

Co-authored-by: Chad Chiang <[email protected]>
Co-authored-by: Gokul Anantha Narayanan <[email protected]>

Author: 3 people

src/sagemaker/image_uris.py

@@ -496,14 +496,7 @@ def _validate_version_and_set_if_needed(version, config, framework, image_scope)
     aliased_versions = list(config.get("version_aliases", {}).keys())
     if len(available_versions) == 1 and version not in aliased_versions:
         return available_versions[0]
-    if not version and framework in [
-        DATA_WRANGLER_FRAMEWORK,
-        HUGGING_FACE_LLM_FRAMEWORK,
-        HUGGING_FACE_TEI_GPU_FRAMEWORK,
-        HUGGING_FACE_TEI_CPU_FRAMEWORK,
-        HUGGING_FACE_LLM_NEURONX_FRAMEWORK,
-        STABILITYAI_FRAMEWORK,
-    ]:
+    if not version:
         version = _get_latest_version(framework, version, image_scope)
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
     return version

tests/unit/sagemaker/serve/detector/test_image_detector.py

@@ -47,7 +47,7 @@ def test_detect_latest_downcast_xgb(self):
         xgb_model = Mock()
         xgb_model.__class__.__bases__ = (xgboost,)
         self.assert_dlc_is_expected(
-            "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.7-1", xgb_model
+            "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1.7-1", xgb_model
         )
 
     @patch.dict(sys.modules, {"xgboost": MagicMock(__version__="1.6.5")})
@@ -57,7 +57,7 @@ def test_detect_earliest_upcast_xgb(self, platform):
         xgb_model = Mock()
         xgb_model.__class__.__bases__ = (xgboost,)
         self.assert_dlc_is_expected(
-            "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.7-1", xgb_model
+            "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1.7-1", xgb_model
         )
 
     @patch.dict(sys.modules, {"xgboost": MagicMock(__version__="1.5.1")})
@@ -67,7 +67,7 @@ def test_detect_exact_match_xgb(self, platform):
         xgb_model = Mock()
         xgb_model.__class__.__bases__ = (xgboost,)
         self.assert_dlc_is_expected(
-            "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1", xgb_model
+            "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1.5-1", xgb_model
         )
 
     # Test Torch Success
@@ -202,7 +202,7 @@ def test_detect_model_no_base(self, platform):
         xgb_model = Mock()
         xgb_model.__class__ = xgboost
         self.assert_dlc_is_expected(
-            "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.7-1", xgb_model
+            "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1.7-1", xgb_model
         )
 
     # helpers