change: update name of jumpstart models accessor, fix small issues

Author: evakravi

src/sagemaker/jumpstart/accessors.py

@@ -34,7 +34,7 @@ def get_sagemaker_version() -> str:
         return SageMakerSettings._parsed_sagemaker_version
 
 
-class JumpStartModelsCache(object):
+class JumpStartModelsAccessor(object):
     """Static class for storing the JumpStart models cache."""
 
     _cache: Optional[cache.JumpStartModelsCache] = None
@@ -67,15 +67,17 @@ def _validate_and_mutate_region_cache_kwargs(
 
     @staticmethod
     def _set_cache_and_region(region: str, cache_kwargs: dict) -> None:
-        """Sets ``JumpStartModelsCache._cache`` and ``JumpStartModelsCache._curr_region``.
+        """Sets ``JumpStartModelsAccessor._cache`` and ``JumpStartModelsAccessor._curr_region``.
 
         Args:
             region (str): region for which to retrieve header/spec.
             cache_kwargs (dict): kwargs to pass to ``JumpStartModelsCache``.
         """
-        if JumpStartModelsCache._cache is None or region != JumpStartModelsCache._curr_region:
-            JumpStartModelsCache._cache = cache.JumpStartModelsCache(region=region, **cache_kwargs)
-            JumpStartModelsCache._curr_region = region
+        if JumpStartModelsAccessor._cache is None or region != JumpStartModelsAccessor._curr_region:
+            JumpStartModelsAccessor._cache = cache.JumpStartModelsCache(
+                region=region, **cache_kwargs
+            )
+            JumpStartModelsAccessor._curr_region = region
 
     @staticmethod
     def get_model_header(region: str, model_id: str, version: str) -> JumpStartModelHeader:
@@ -86,12 +88,12 @@ def get_model_header(region: str, model_id: str, version: str) -> JumpStartModel
             model_id (str): model id to retrieve.
             version (str): semantic version to retrieve for the model id.
         """
-        cache_kwargs = JumpStartModelsCache._validate_and_mutate_region_cache_kwargs(
-            JumpStartModelsCache._cache_kwargs, region
+        cache_kwargs = JumpStartModelsAccessor._validate_and_mutate_region_cache_kwargs(
+            JumpStartModelsAccessor._cache_kwargs, region
         )
-        JumpStartModelsCache._set_cache_and_region(region, cache_kwargs)
-        assert JumpStartModelsCache._cache is not None
-        return JumpStartModelsCache._cache.get_header(model_id, version)
+        JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
+        assert JumpStartModelsAccessor._cache is not None
+        return JumpStartModelsAccessor._cache.get_header(model_id, version)
 
     @staticmethod
     def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelSpecs:
@@ -102,12 +104,12 @@ def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelS
             model_id (str): model id to retrieve.
             version (str): semantic version to retrieve for the model id.
         """
-        cache_kwargs = JumpStartModelsCache._validate_and_mutate_region_cache_kwargs(
-            JumpStartModelsCache._cache_kwargs, region
+        cache_kwargs = JumpStartModelsAccessor._validate_and_mutate_region_cache_kwargs(
+            JumpStartModelsAccessor._cache_kwargs, region
         )
-        JumpStartModelsCache._set_cache_and_region(region, cache_kwargs)
-        assert JumpStartModelsCache._cache is not None
-        return JumpStartModelsCache._cache.get_specs(model_id, version)
+        JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
+        assert JumpStartModelsAccessor._cache is not None
+        return JumpStartModelsAccessor._cache.get_specs(model_id, version)
 
     @staticmethod
     def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None:
@@ -120,18 +122,18 @@ def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None:
             cache_kwargs (str): cache kwargs to validate.
             region (str): Optional. The region to validate along with the kwargs.
         """
-        cache_kwargs = JumpStartModelsCache._validate_and_mutate_region_cache_kwargs(
+        cache_kwargs = JumpStartModelsAccessor._validate_and_mutate_region_cache_kwargs(
             cache_kwargs, region
         )
-        JumpStartModelsCache._cache_kwargs = cache_kwargs
+        JumpStartModelsAccessor._cache_kwargs = cache_kwargs
         if region is None:
-            JumpStartModelsCache._cache = cache.JumpStartModelsCache(
-                **JumpStartModelsCache._cache_kwargs
+            JumpStartModelsAccessor._cache = cache.JumpStartModelsCache(
+                **JumpStartModelsAccessor._cache_kwargs
             )
         else:
-            JumpStartModelsCache._curr_region = region
-            JumpStartModelsCache._cache = cache.JumpStartModelsCache(
-                region=region, **JumpStartModelsCache._cache_kwargs
+            JumpStartModelsAccessor._curr_region = region
+            JumpStartModelsAccessor._cache = cache.JumpStartModelsCache(
+                region=region, **JumpStartModelsAccessor._cache_kwargs
             )
 
     @staticmethod
@@ -146,4 +148,4 @@ def reset_cache(cache_kwargs: Dict[str, Any] = None, region: Optional[str] = Non
             region (str): The region to validate along with the kwargs.
         """
         cache_kwargs_dict = {} if cache_kwargs is None else cache_kwargs
-        JumpStartModelsCache.set_cache_kwargs(cache_kwargs_dict, region)
+        JumpStartModelsAccessor.set_cache_kwargs(cache_kwargs_dict, region)

src/sagemaker/jumpstart/artifacts.py

@@ -49,7 +49,7 @@ def _retrieve_image_uri(
     Args:
         model_id (str): JumpStart model ID for which to retrieve image URI.
         model_version (str): Version of the JumpStart model for which to retrieve
-            the image URI (default: None).
+            the image URI.
         image_scope (str): The image type, i.e. what it is used for.
             Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
             ``image_scope`` is ignored.
@@ -67,12 +67,10 @@ def _retrieve_image_uri(
         container_version (str): the version of docker image.
             Ideally the value of parameter should be created inside the framework.
             For custom use, see the list of supported container versions:
-            https://github.com/aws/deep-learning-containers/blob/master/available_images.md
-            (default: None).
+            https://github.com/aws/deep-learning-containers/blob/master/available_images.md.
         distribution (dict): A dictionary with information on how to run distributed training
         training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
-            A configuration class for the SageMaker Training Compiler
-            (default: None).
+            A configuration class for the SageMaker Training Compiler.
 
     Returns:
         str: the ECR URI for the corresponding SageMaker Docker image.
@@ -94,7 +92,7 @@ def _retrieve_image_uri(
             f"JumpStart models only support scopes: {', '.join(SUPPORTED_JUMPSTART_SCOPES)}."
         )
 
-    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+    model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
         region, model_id, model_version
     )
 
@@ -202,7 +200,7 @@ def _retrieve_model_uri(
             f"JumpStart models only support scopes: {', '.join(SUPPORTED_JUMPSTART_SCOPES)}."
         )
 
-    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+    model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
         region, model_id, model_version
     )
     if model_scope == INFERENCE:
@@ -261,7 +259,7 @@ def _retrieve_script_uri(
             f"JumpStart models only support scopes: {', '.join(SUPPORTED_JUMPSTART_SCOPES)}."
         )
 
-    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+    model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
         region, model_id, model_version
     )
     if script_scope == INFERENCE:

tests/unit/sagemaker/image_uris/jumpstart/test_catboost.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_catboost_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_catboost_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_common.py

@@ -21,7 +21,7 @@
 from sagemaker.jumpstart import constants as sagemaker_constants
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_common_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_spec_from_base_spec

tests/unit/sagemaker/image_uris/jumpstart/test_huggingface.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_huggingface_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_lightgbm.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_lightgbm_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_lightgbm_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_mxnet.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_mxnet_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_mxnet_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_pytorch.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_pytorch_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_pytorch_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_sklearn.py

@@ -23,7 +23,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sklearn_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -32,7 +32,7 @@ def test_jumpstart_sklearn_image_uri(patched_get_model_specs):
     instance_type = "ml.m2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_tensorflow.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_tensorflow_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_tensorflow_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(

tests/unit/sagemaker/image_uris/jumpstart/test_xgboost.py

@@ -22,7 +22,7 @@
 from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 
-@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_xgboost_image_uri(patched_get_model_specs):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
@@ -31,7 +31,7 @@ def test_jumpstart_xgboost_image_uri(patched_get_model_specs):
     instance_type = "ml.p2.xlarge"
     region = "us-west-2"
 
-    model_specs = accessors.JumpStartModelsCache.get_model_specs(region, model_id, model_version)
+    model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
 
     # inference
     uri = image_uris.retrieve(