fix: Add regex for short-form sagemaker-xgboost tags in framework_name_from_image

Author: Mark Bunday

src/sagemaker/fw_utils.py

@@ -410,6 +410,8 @@ def framework_name_from_image(image_uri):
             '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-rl-<fw>:<rl_toolkit><rl_version>-<device>-<py_ver>'
             current:
             '<account>.dkr.ecr.<region>.amazonaws.com/<fw>-<image_scope>:<fw_version>-<device>-<py_ver>'
+            current:
+            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-xgboost:<fw_version>-<container_version>'
 
     Returns:
         tuple: A tuple containing:
@@ -450,6 +452,16 @@ def framework_name_from_image(image_uri):
     legacy_match = legacy_name_pattern.match(sagemaker_match.group(9))
     if legacy_match is not None:
         return (legacy_match.group(1), legacy_match.group(2), legacy_match.group(4), None)
+
+    # sagemaker-xgboost images are tagged with two aliases, e.g.:
+    # 1. Long tag: "315553699071.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1-cpu-py3"
+    # 2. Short tag: "315553699071.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1"
+    # Note 1: Both tags point to the same image
+    # Note 2: Both tags have full GPU capabilities, despite "cpu" delineation in the long tag
+    short_xgboost_tag_pattern = re.compile(r"^sagemaker-(xgboost):(.*)$")
+    short_xgboost_tag_match = short_xgboost_tag_pattern.match(sagemaker_match.group(9))
+    if short_xgboost_tag_match is not None:
+        return (short_xgboost_tag_match.group(1), "py3", short_xgboost_tag_match.group(2), None)
     return None, None, None, None
 
 

tests/unit/test_fw_utils.py

@@ -511,6 +511,39 @@ def test_framework_version_from_tag_other():
     assert version is None
 
 
+def test_framework_name_from_xgboost_image_short_tag():
+    ecr_uri = "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost"
+    image_tag = "1.5-1"
+    image_uri = f"{ecr_uri}:{image_tag}"
+    expected_result = ("xgboost", "py3", "1.5-1", None)
+    assert expected_result == fw_utils.framework_name_from_image(image_uri)
+
+
+def test_framework_name_from_xgboost_image_long_tag():
+    ecr_uri = "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost"
+    image_tag = "1.5-1-cpu-py3"
+    image_uri = f"{ecr_uri}:{image_tag}"
+    expected_result = ("xgboost", "py3", "1.5-1-cpu-py3", None)
+    assert expected_result == fw_utils.framework_name_from_image(image_uri)
+
+
+def test_framework_name_from_xgboost_image_short_long_tag_equal():
+    ecr_uri = "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost"
+    short_image_tag = "1.5-1"
+    long_image_tag = "1.5-1-cpu-py3"
+    short_image_uri = f"{ecr_uri}:{short_image_tag}"
+    long_image_uri = f"{ecr_uri}:{long_image_tag}"
+    expected_framework = "xgboost"
+    expected_py = "py3"
+
+    short_framework, short_py, short_tag, _ = fw_utils.framework_name_from_image(short_image_uri)
+    long_framework, long_py, long_tag, _ = fw_utils.framework_name_from_image(long_image_uri)
+    assert expected_framework == short_framework == long_framework
+    assert expected_py == short_py == long_py
+    assert short_image_tag == short_tag
+    assert long_image_tag == long_tag
+
+
 def test_model_code_key_prefix_with_all_values_present():
     key_prefix = fw_utils.model_code_key_prefix("prefix", "model_name", "image_uri")
     assert key_prefix == "prefix/model_name"