change: improve getting sm version logic

Author: evakravi

src/sagemaker/jumpstart/constants.py

@@ -27,3 +27,5 @@
 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

@@ -75,6 +75,17 @@ def get_formatted_manifest(
 def get_sagemaker_version() -> str:
     """Returns sagemaker library version.
 
+    If the sagemaker library version has not been set yet, this function
+    calls ``parse_sagemaker_version`` to retrive the version.
+    """
+    if constants.PARSED_SAGEMAKER_VERSION == "":
+        constants.PARSED_SAGEMAKER_VERSION = parse_sagemaker_version()
+    return constants.PARSED_SAGEMAKER_VERSION
+
+
+def parse_sagemaker_version() -> str:
+    """Returns sagemaker library version. This should only be called once.
+
     Function reads ``__version__`` variable in ``sagemaker`` module.
     In order to maintain compatibility with the ``semantic_version``
     library, versions with fewer than 2, or more than 3, periods are rejected.

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
-from mock.mock import patch
+from mock.mock import Mock, patch
 import pytest
 from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.constants import JUMPSTART_REGION_NAME_SET
@@ -73,33 +73,42 @@ def test_get_formatted_manifest():
     assert utils.get_formatted_manifest([]) == {}
 
 
-def test_get_sagemaker_version():
+def test_parse_sagemaker_version():
 
     with patch("sagemaker.__version__", "1.2.3"):
-        assert utils.get_sagemaker_version() == "1.2.3"
+        assert utils.parse_sagemaker_version() == "1.2.3"
 
     with patch("sagemaker.__version__", "1.2.3.3332j"):
-        assert utils.get_sagemaker_version() == "1.2.3"
+        assert utils.parse_sagemaker_version() == "1.2.3"
 
     with patch("sagemaker.__version__", "1.2.3."):
-        assert utils.get_sagemaker_version() == "1.2.3"
+        assert utils.parse_sagemaker_version() == "1.2.3"
 
     with pytest.raises(ValueError):
         with patch("sagemaker.__version__", "1.2.3dfsdfs"):
-            utils.get_sagemaker_version()
+            utils.parse_sagemaker_version()
 
     with pytest.raises(RuntimeError):
         with patch("sagemaker.__version__", "1.2"):
-            utils.get_sagemaker_version()
+            utils.parse_sagemaker_version()
 
     with pytest.raises(RuntimeError):
         with patch("sagemaker.__version__", "1"):
-            utils.get_sagemaker_version()
+            utils.parse_sagemaker_version()
 
     with pytest.raises(RuntimeError):
         with patch("sagemaker.__version__", ""):
-            utils.get_sagemaker_version()
+            utils.parse_sagemaker_version()
 
     with pytest.raises(RuntimeError):
         with patch("sagemaker.__version__", "1.2.3.4.5"):
-            utils.get_sagemaker_version()
+            utils.parse_sagemaker_version()
+
+
+@patch("sagemaker.jumpstart.utils.parse_sagemaker_version")
+@patch("sagemaker.jumpstart.constants.PARSED_SAGEMAKER_VERSION", "")
+def test_get_sagemaker_version(patched_parse_sm_version: Mock):
+    utils.get_sagemaker_version()
+    utils.get_sagemaker_version()
+    utils.get_sagemaker_version()
+    assert patched_parse_sm_version.called_only_once()