fix image retriever tests (#5831)

* fix image retriever tests

* fix: skip test_retrieve_image_uri_intelligent_default for now

* fix: bump image uri assertion in test_image_retriever

* fix: add missing import

---------

Co-authored-by: Syed Jafri <syedjfr@amazon.com>

Author: mujtaba1747

sagemaker-core/src/sagemaker/core/image_retriever/image_retriever.py

@@ -52,6 +52,8 @@
 
 
 class ImageRetriever:
+    _config = SageMakerConfig()
+
     @staticmethod
     def retrieve_hugging_face_uri(
         region: str,
@@ -110,7 +112,7 @@ def retrieve_hugging_face_uri(
         args = dict(locals())
         for name, val in args.items():
             if name in CONFIGURABLE_ATTRIBUTES and not val:
-                default_value = SageMakerConfig.resolve_value_from_config(
+                default_value = ImageRetriever._config.resolve_value_from_config(
                     config_path=_simple_path(
                         SAGEMAKER, MODULES, IMAGE_RETRIEVER, to_camel_case(name)
                     )
@@ -499,7 +501,7 @@ def retrieve(
         args = dict(locals())
         for name, val in args.items():
             if name in CONFIGURABLE_ATTRIBUTES and not val:
-                default_value = SageMakerConfig.resolve_value_from_config(
+                default_value = ImageRetriever._config.resolve_value_from_config(
                     config_path=_simple_path(
                         SAGEMAKER, MODULES, IMAGE_RETRIEVER, to_camel_case(name)
                     )

sagemaker-core/src/sagemaker/core/image_retriever/image_retriever_utils.py

@@ -28,6 +28,7 @@
     GRAVITON_ALLOWED_FRAMEWORKS,
 )
 from sagemaker.core.common_utils import _botocore_resolver, get_instance_type_family
+from sagemaker.core.spark import defaults
 
 logger = logging.getLogger(__name__)
 

sagemaker-core/tests/integ/image_retriever/test_image_retriever.py

@@ -12,7 +12,6 @@
 from sagemaker.core.config.config_manager import SageMakerConfig
 
 
-@pytest.mark.skip("Disabling this for now, Need to be fixed")
 @pytest.mark.integ
 def test_retrieve_image_uri():
     image_uri = ImageRetriever.retrieve("clarify", "us-west-2")
@@ -28,7 +27,7 @@ def test_retrieve_image_uri():
     )
     assert (
         image_uri
-        == "053634841547.dkr.ecr.us-west-1.amazonaws.com/sagemaker-distribution-prod:3.0.0-gpu"
+        == "053634841547.dkr.ecr.us-west-1.amazonaws.com/sagemaker-distribution-prod:3.2.0-gpu"
     )
 
     image_uri = ImageRetriever.retrieve(
@@ -56,7 +55,6 @@ def test_retrieve_image_uri():
     )
 
 
-@pytest.mark.skip("Disabling this for now, Need to be fixed")
 @pytest.mark.integ
 def test_retrieve_pytorch_uri():
     image_uri = ImageRetriever.retrieve_pytorch_uri(
@@ -72,7 +70,6 @@ def test_retrieve_pytorch_uri():
     )
 
 
-@pytest.mark.skip("Disabling this for now, Need to be fixed")
 @pytest.mark.integ
 def test_retrieve_hugging_face_uri():
     image_uri = ImageRetriever.retrieve_hugging_face_uri(
@@ -84,22 +81,23 @@ def test_retrieve_hugging_face_uri():
         base_framework_version="pytorch2.0.0",
         container_version="cu110-ubuntu20.04",
     )
-    assert image_uri == "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training"
-    ":2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
+    assert (
+        image_uri
+        == "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training"
+        ":2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
+    )
 
 
-@pytest.mark.skip("Disabling this for now, Need to be fixed")
 @pytest.mark.integ
 def test_retrieve_base_python_image_uri():
-    image_uri = ImageRetriever.retrieve_base_python_image_uri()
+    image_uri = ImageRetriever.retrieve_base_python_image_uri(region="us-west-2")
     assert image_uri == "236514542706.dkr.ecr.us-west-2.amazonaws.com/sagemaker-base-python-310:1.0"
 
 
-@pytest.mark.skip("Disabling this for now, Need to be fixed")
-@pytest.mark.integ
+@pytest.mark.skip(reason="Test is failing due to locals()[name] = default_value in Image Retriever")
 @patch.object(SageMakerConfig, "resolve_value_from_config")
 def test_retrieve_image_uri_intelligent_default(mock_load_config):
-    def custom_return(config_path):
+    def custom_return(config_path=None, **kwargs):
         if config_path == _simple_path(
             SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, "ImageScope"
         ):