handle invalid sev ver and incompatible sagemaker versions

Author: Eli Davidson

src/sagemaker/jumpstart/cache.py

@@ -262,7 +262,7 @@ def _model_id_retrieval_function(
             return JumpStartVersionedModelId(model_id, sm_compatible_model_version)
 
         versions_incompatible_with_sagemaker = [
-            Version(header.version)
+            header.version
             for header in manifest.values()  # type: ignore
             if header.model_id == model_id
         ]

src/sagemaker/jumpstart/utils.py

@@ -20,7 +20,7 @@
 from urllib.parse import urlparse
 import boto3
 from botocore.exceptions import ClientError
-from packaging.version import Version
+from packaging.version import Version, InvalidVersion
 import botocore
 import sagemaker
 from sagemaker.config.config_schema import (
@@ -1527,4 +1527,6 @@ def wrapped_f(*args, **kwargs):
     return wrapper_cache(_func)
 
 def get_latest_version(versions: List[str]) -> Optional[str]:
-    return None if not versions else max(versions, key=Version)
+    try: return None if not versions else max(versions, key=Version)
+    except InvalidVersion as e:
+        return max(versions)

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -22,7 +22,10 @@
 from mock.mock import MagicMock
 import pytest
 from mock import patch
+from packaging.version import Version
 
+
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.cache import (
     JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
     JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY,
@@ -33,6 +36,7 @@
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
 )
 from sagemaker.jumpstart.types import (
+    JumpStartCachedContentValue,
     JumpStartModelHeader,
     JumpStartModelSpecs,
     JumpStartVersionedModelId,
@@ -50,7 +54,7 @@
     BASE_PROPRIETARY_SPEC,
     BASE_PROPRIETARY_MANIFEST,
 )
-from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.jumpstart.utils import get_formatted_manifest, get_jumpstart_content_bucket
 
 
 @patch.object(JumpStartModelsCache, "_retrieval_function", patched_retrieval_function)
@@ -1119,3 +1123,130 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
             ),
         ]
     )
+
+@patch.object(JumpStartModelsCache, "_retrieval_function")
+def test_jumpstart_cache_handles_versioning_correctly_for_open_source_weights(
+   retrieval_function: Mock
+):
+    sm_version = Version(utils.get_sagemaker_version())
+    new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
+    print(str(new_sm_version))
+    versions = ["1.0.0", "2.9.1", "2.16.0"]
+    manifest = [
+        {
+            "model_id": "test-model",
+            "version": version,
+            "min_version": "2.49.0",
+            "spec_key": "spec_key"
+        }
+        for version in versions
+    ]
+
+    manifest.append(
+        {
+            "model_id": "test-model",
+            "version": "3.0.0",
+            "min_version": str(new_sm_version),
+            "spec_key": "spec_key"
+        }
+    )
+    
+    manifest_dict = {}
+    for header in manifest:
+        header_obj = JumpStartModelHeader(header)
+        manifest_dict[JumpStartVersionedModelId(header_obj.model_id, header_obj.version)] = (
+            header_obj
+        )
+    retrieval_function.return_value = JumpStartCachedContentValue(
+        formatted_content=manifest_dict
+    )
+    key = JumpStartVersionedModelId("test-model", "*")
+
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    result = cache._get_open_weight_manifest_key_from_model_id( key = key, value = None )
+
+    assert_key = JumpStartVersionedModelId("test-model", "2.16.0")
+
+    assert result == assert_key
+
+@patch.object(JumpStartModelsCache, "_retrieval_function")
+def test_jumpstart_cache_handles_versioning_correctly_for_proprietary_weights(
+   retrieval_function: Mock
+):
+    sm_version = Version(utils.get_sagemaker_version())
+    new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
+    print(str(new_sm_version))
+    versions = ["1.0.0", "2.9.1", "2.16.0"]
+    manifest = [
+        {
+            "model_id": "test-model",
+            "version": version,
+            "min_version": "2.49.0",
+            "spec_key": "spec_key"
+        }
+        for version in versions
+    ]
+
+    manifest.append(
+        {
+            "model_id": "test-model",
+            "version": "3.0.0",
+            "min_version": str(new_sm_version),
+            "spec_key": "spec_key"
+        }
+    )
+    
+    manifest_dict = {}
+    for header in manifest:
+        header_obj = JumpStartModelHeader(header)
+        manifest_dict[JumpStartVersionedModelId(header_obj.model_id, header_obj.version)] = (
+            header_obj
+        )
+    retrieval_function.return_value = JumpStartCachedContentValue(
+        formatted_content=manifest_dict
+    )
+    key = JumpStartVersionedModelId("test-model", "*")
+
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    result = cache._get_proprietary_manifest_key_from_model_id( key = key, value = None )
+
+    assert_key = JumpStartVersionedModelId("test-model", "2.16.0")
+
+    assert result == assert_key
+
+
+@patch.object(JumpStartModelsCache, "_retrieval_function")
+def test_jumpstart_cache_handles_versioning_correctly_non_sem_ver(
+   retrieval_function: Mock
+):
+    sm_version = Version(utils.get_sagemaker_version())
+    new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
+    print(str(new_sm_version))
+    versions = ["abc", "2.9.1", "2.16.0"]
+    manifest = [
+        {
+            "model_id": "test-model",
+            "version": version,
+            "min_version": "2.49.0",
+            "spec_key": "spec_key"
+        }
+        for version in versions
+    ]
+    
+    manifest_dict = {}
+    for header in manifest:
+        header_obj = JumpStartModelHeader(header)
+        manifest_dict[JumpStartVersionedModelId(header_obj.model_id, header_obj.version)] = (
+            header_obj
+        )
+    retrieval_function.return_value = JumpStartCachedContentValue(
+        formatted_content=manifest_dict
+    )
+    key = JumpStartVersionedModelId("test-model", "*")
+
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    result = cache._get_open_weight_manifest_key_from_model_id( key = key, value = None )
+
+    assert_key = JumpStartVersionedModelId("test-model", "abc")
+
+    assert result == assert_key

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -2151,6 +2151,9 @@ def test_get_latest_version_empty_list_is_none():
 def test_get_latest_version_none_is_none():
     assert utils.get_latest_version(None) == None
 
+def test_get_latest_version_with_invalid_sem_ver():
+    assert utils.get_latest_version(["2.9.1", "2.16.0", "1.0.0", "abc"]) == "abc"
+
 @pytest.mark.parametrize(
     "data, expected",
     [(None, []), ([], []), (get_base_deployment_configs_metadata(), get_base_deployment_configs())],