fix: sagemaker version for jumpstart

Author: evakravi

src/sagemaker/jumpstart/constants.py

@@ -27,5 +27,3 @@
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
-
-PARSED_SAGEMAKER_VERSION = ""  # this gets set by sagemaker.jumpstart.utils.get_sagemaker_version()

src/sagemaker/jumpstart/utils.py

@@ -19,6 +19,22 @@
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartVersionedModelId
 
 
+class SageMakerSettings(object):
+    """Static class for storing the SageMaker settings."""
+
+    _PARSED_SAGEMAKER_VERSION = ""
+
+    @staticmethod
+    def set_sagemaker_version(version: str) -> None:
+        """Set SageMaker version."""
+        SageMakerSettings._PARSED_SAGEMAKER_VERSION = version
+
+    @staticmethod
+    def get_sagemaker_version() -> str:
+        """Return SageMaker version."""
+        return SageMakerSettings._PARSED_SAGEMAKER_VERSION
+
+
 def get_jumpstart_launched_regions_message() -> str:
     """Returns formatted string indicating where JumpStart is launched."""
     if len(constants.JUMPSTART_REGION_NAME_SET) == 0:
@@ -79,9 +95,9 @@ def get_sagemaker_version() -> str:
     calls ``parse_sagemaker_version`` to retrieve the version and set
     the constant.
     """
-    if constants.PARSED_SAGEMAKER_VERSION == "":
-        constants.PARSED_SAGEMAKER_VERSION = parse_sagemaker_version()
-    return constants.PARSED_SAGEMAKER_VERSION
+    if SageMakerSettings.get_sagemaker_version() == "":
+        SageMakerSettings.set_sagemaker_version(parse_sagemaker_version())
+    return SageMakerSettings.get_sagemaker_version()
 
 
 def parse_sagemaker_version() -> str:

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -106,7 +106,7 @@ def test_parse_sagemaker_version():
 
 
 @patch("sagemaker.jumpstart.utils.parse_sagemaker_version")
-@patch("sagemaker.jumpstart.constants.PARSED_SAGEMAKER_VERSION", "")
+@patch("sagemaker.jumpstart.utils.SageMakerSettings._PARSED_SAGEMAKER_VERSION", "")
 def test_get_sagemaker_version(patched_parse_sm_version: Mock):
     utils.get_sagemaker_version()
     utils.get_sagemaker_version()