diff --git a/src/sagemaker/jumpstart/cache.py b/src/sagemaker/jumpstart/cache.py index 8ac813a6c5..bdfc01cba3 100644 --- a/src/sagemaker/jumpstart/cache.py +++ b/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 ] @@ -540,9 +540,7 @@ def _select_version( """ if version_str == "*": - if len(available_versions) == 0: - return None - return str(max(available_versions)) + return utils.get_latest_version(available_versions) if model_type == JumpStartModelType.PROPRIETARY: if "*" in version_str: diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py index 46e5f8a847..23245b24e5 100644 --- a/src/sagemaker/jumpstart/utils.py +++ b/src/sagemaker/jumpstart/utils.py @@ -21,7 +21,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 from sagemaker_core.shapes import ModelAccessConfig import sagemaker @@ -1630,3 +1630,11 @@ def get_draft_model_content_bucket(provider: Dict, region: str) -> str: return get_jumpstart_gated_content_bucket(region=region) return get_jumpstart_content_bucket(region=region) return neo_bucket + + +def get_latest_version(versions: List[str]) -> Optional[str]: + """Returns the latest version using sem-ver when possible.""" + try: + return None if not versions else max(versions, key=Version) + except InvalidVersion: + return max(versions) diff --git a/tests/unit/sagemaker/jumpstart/test_cache.py b/tests/unit/sagemaker/jumpstart/test_cache.py index da20debc6a..6816983542 100644 --- a/tests/unit/sagemaker/jumpstart/test_cache.py +++ b/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, @@ -1119,3 +1123,124 @@ 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 diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py index 7cf8fdc9b6..ea4d64f289 100644 --- a/tests/unit/sagemaker/jumpstart/test_utils.py +++ b/tests/unit/sagemaker/jumpstart/test_utils.py @@ -2144,6 +2144,22 @@ def test_has_instance_rate_stat(stats, expected): assert utils.has_instance_rate_stat(stats) is expected +def test_get_latest_version(): + assert utils.get_latest_version(["2.9.1", "2.16.0", "1.0.0"]) == "2.16.0" + + +def test_get_latest_version_empty_list_is_none(): + assert utils.get_latest_version([]) is None + + +def test_get_latest_version_none_is_none(): + assert utils.get_latest_version(None) is 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())],