Skip to content

Commit 4ff8e11

Browse files
giangle-AIMohamed Zeidan
andauthored
fix: resolve HuggingFace version alias in image URI tag (#5722)
* fix: resolve HuggingFace version alias in image URI tag When constructing the image URI tag for HuggingFace framework containers, the code used `original_version` (the raw user input, e.g. "4.49") instead of resolving it through the version alias mapping (e.g. "4.49" → "4.49.0"). This caused image URIs to contain truncated transformer versions like `transformers4.49` instead of the correct `transformers4.49.0`, resulting in references to non-existent container images. Use `_version_for_config(version, config)` to properly resolve the version alias before embedding it in the image tag. Made-with: Cursor * test: add comprehensive version alias tests for HuggingFace image URIs - Fix existing test_retrieve_huggingface assertions to expect resolved versions (4.2.1) instead of raw aliases (4.2) in image tags - Add parametrized test covering all 25 version aliases across training and inference scopes to verify aliases resolve correctly in tags - Add parametrized test covering all 26 non-aliased (full) versions to confirm no regression when passing complete version strings directly Made-with: Cursor * style: apply black formatting to HuggingFace alias tests The new parametrized HuggingFace version-alias tests were not formatted with the repo's pinned black (26.3.1), failing the codestyle-doc-tests (sagemaker-core) check. Formatting only - no behavior change. --------- Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>
1 parent 6714174 commit 4ff8e11

2 files changed

Lines changed: 83 additions & 3 deletions

File tree

‎sagemaker-core/src/sagemaker/core/image_uris.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -258,7 +258,7 @@ def retrieve(
258258
pt_or_tf_version = (
259259
re.compile("^(pytorch|tensorflow)(.*)$").match(base_framework_version).group(2)
260260
)
261-
_version = original_version
261+
_version = _version_for_config(version, config)
262262

263263
if repo in [
264264
"huggingface-pytorch-trcomp-training",

‎sagemaker-core/tests/unit/image_uris/test_retrieve.py‎

Lines changed: 82 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -715,7 +715,7 @@ def test_retrieve_huggingface(config_for_framework):
715715
)
716716
assert (
717717
"564829616587.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:"
718-
"1.6-transformers4.2-gpu-py37-cu110-ubuntu18.04" == pt_uri_mv
718+
"1.6-transformers4.2.1-gpu-py37-cu110-ubuntu18.04" == pt_uri_mv
719719
)
720720

721721
pt_uri = image_uris.retrieve(
@@ -760,7 +760,7 @@ def test_retrieve_huggingface(config_for_framework):
760760
)
761761
assert (
762762
"564829616587.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:"
763-
"1.6.0-transformers4.3.1-gpu-py37-cu110-ubuntu18.04" == pt_new_version
763+
"1.6.0-transformers4.2.1-gpu-py37-cu110-ubuntu18.04" == pt_new_version
764764
)
765765

766766

@@ -832,6 +832,86 @@ def test_get_latest_version_function_with_no_framework(config_for_framework):
832832
assert "No framework config for framework" in str(e.exception)
833833

834834

835+
def _get_huggingface_alias_test_cases():
836+
"""Build parametrized test cases for every HuggingFace version alias."""
837+
config = image_uris.config_for_framework("huggingface")
838+
cases = []
839+
for scope in ("training", "inference"):
840+
section = config[scope]
841+
aliases = section.get("version_aliases", {})
842+
for alias, resolved in aliases.items():
843+
ver_cfg = section["versions"][resolved]
844+
base_fws = [k for k in ver_cfg if k != "version_aliases"]
845+
base_fw = base_fws[0]
846+
py_ver = ver_cfg[base_fw]["py_versions"][0]
847+
inst = "ml.p3.2xlarge" if scope == "training" else "ml.c5.xlarge"
848+
cases.append(
849+
pytest.param(
850+
scope, alias, resolved, base_fw, py_ver, inst, id=f"{scope}-{alias}->{resolved}"
851+
)
852+
)
853+
return cases
854+
855+
856+
def _get_huggingface_full_version_test_cases():
857+
"""Build parametrized test cases for every non-aliased HuggingFace version."""
858+
config = image_uris.config_for_framework("huggingface")
859+
cases = []
860+
for scope in ("training", "inference"):
861+
section = config[scope]
862+
for full_ver, ver_cfg in section["versions"].items():
863+
base_fws = [k for k in ver_cfg if k != "version_aliases"]
864+
base_fw = base_fws[0]
865+
py_ver = ver_cfg[base_fw]["py_versions"][0]
866+
inst = "ml.p3.2xlarge" if scope == "training" else "ml.c5.xlarge"
867+
cases.append(
868+
pytest.param(scope, full_ver, base_fw, py_ver, inst, id=f"{scope}-{full_ver}")
869+
)
870+
return cases
871+
872+
873+
@pytest.mark.parametrize(
874+
"scope,alias,resolved,base_fw,py_ver,instance_type",
875+
_get_huggingface_alias_test_cases(),
876+
)
877+
def test_huggingface_version_alias_resolves_in_tag(
878+
scope, alias, resolved, base_fw, py_ver, instance_type
879+
):
880+
"""Version aliases must be resolved to full versions in image URI tags."""
881+
uri = image_uris.retrieve(
882+
framework="huggingface",
883+
region="us-east-1",
884+
version=alias,
885+
py_version=py_ver,
886+
image_scope=scope,
887+
base_framework_version=base_fw,
888+
instance_type=instance_type,
889+
)
890+
assert (
891+
f"transformers{resolved}-" in uri
892+
), f"Expected resolved version 'transformers{resolved}-' in URI, got: {uri}"
893+
894+
895+
@pytest.mark.parametrize(
896+
"scope,full_version,base_fw,py_ver,instance_type",
897+
_get_huggingface_full_version_test_cases(),
898+
)
899+
def test_huggingface_full_version_in_tag(scope, full_version, base_fw, py_ver, instance_type):
900+
"""Full (non-aliased) versions must appear unchanged in image URI tags."""
901+
uri = image_uris.retrieve(
902+
framework="huggingface",
903+
region="us-east-1",
904+
version=full_version,
905+
py_version=py_ver,
906+
image_scope=scope,
907+
base_framework_version=base_fw,
908+
instance_type=instance_type,
909+
)
910+
assert (
911+
f"transformers{full_version}-" in uri
912+
), f"Expected full version 'transformers{full_version}-' in URI, got: {uri}"
913+
914+
835915
@pytest.mark.parametrize(
836916
"framework",
837917
[

0 commit comments

Comments
 (0)