Skip to content

Commit 305bdf8

Browse files
authored
chore: fix semantic versioning for wildcard identifier (aws#5105)
1 parent 230fb55 commit 305bdf8

File tree

3 files changed

+42
-0
lines changed

3 files changed

+42
-0
lines changed

src/sagemaker/jumpstart/cache.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -552,6 +552,12 @@ def _select_version(
552552
)
553553
return version_str if version_str in available_versions else None
554554

555+
if version_str[-1] == "*":
556+
# major or minor version is pinned, e.g 1.* or 1.0.*
557+
return utils.get_latest_version(
558+
[version for version in available_versions if version.startswith(version_str[:-1])]
559+
)
560+
555561
try:
556562
spec = SpecifierSet(f"=={version_str}")
557563
except InvalidSpecifier:

tests/unit/sagemaker/jumpstart/constants.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15990,6 +15990,18 @@
1599015990
"spec_key": "community_models_specs/tensorflow-ic-"
1599115991
"imagenet-inception-v3-classification-4/specs_v3.0.0.json",
1599215992
},
15993+
{
15994+
"model_id": "meta-textgeneration-llama-2-7b",
15995+
"version": "4.9.0",
15996+
"min_version": "2.49.0",
15997+
"spec_key": "community_models/meta-textgeneration-llama-2-7b/specs_v4.9.0.json",
15998+
},
15999+
{
16000+
"model_id": "meta-textgeneration-llama-2-7b",
16001+
"version": "4.13.0",
16002+
"min_version": "2.49.0",
16003+
"spec_key": "community_models/meta-textgeneration-llama-2-7b/specs_v4.13.0.json",
16004+
},
1599316005
]
1599416006

1599516007
BASE_PROPRIETARY_HEADER = {

tests/unit/sagemaker/jumpstart/test_cache.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,30 @@ def test_jumpstart_cache_get_header():
184184
semantic_version_str="1.0.*",
185185
)
186186

187+
assert JumpStartModelHeader(
188+
{
189+
"model_id": "meta-textgeneration-llama-2-7b",
190+
"version": "4.13.0",
191+
"min_version": "2.49.0",
192+
"spec_key": "community_models/meta-textgeneration-llama-2-7b/specs_v4.13.0.json",
193+
}
194+
) == cache.get_header(
195+
model_id="meta-textgeneration-llama-2-7b",
196+
semantic_version_str="*",
197+
)
198+
199+
assert JumpStartModelHeader(
200+
{
201+
"model_id": "meta-textgeneration-llama-2-7b",
202+
"version": "4.13.0",
203+
"min_version": "2.49.0",
204+
"spec_key": "community_models/meta-textgeneration-llama-2-7b/specs_v4.13.0.json",
205+
}
206+
) == cache.get_header(
207+
model_id="meta-textgeneration-llama-2-7b",
208+
semantic_version_str="4.*",
209+
)
210+
187211
assert JumpStartModelHeader(
188212
{
189213
"model_id": "ai21-summarization",

0 commit comments

Comments
 (0)