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

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
 
 
@@ -459,12 +471,16 @@ def framework_version_from_tag(image_tag):
     Args:
         image_tag (str): Image tag, which should take the form
             '<framework_version>-<device>-<py_version>'
+            '<xgboost_version>-<container_version>'
 
     Returns:
         str: The framework version.
     """
     tag_pattern = re.compile(r"^(.*)-(cpu|gpu)-(py2|py3\d*)$")
     tag_match = tag_pattern.match(image_tag)
+    if tag_match is None:
+        short_xgboost_tag_pattern = re.compile(r"^(\d\.\d+\-\d)$")
+        tag_match = short_xgboost_tag_pattern.match(image_tag)
     return None if tag_match is None else tag_match.group(1)
 
 

tests/unit/test_fw_utils.py

@@ -511,6 +511,33 @@ def test_framework_version_from_tag_other():
     assert version is None
 
 
+def test_xgboost_version_from_tag():
+    tags = (
+        "1.5-1-cpu-py3",
+        "1.5-1",
+    )
+
+    for tag in tags:
+        version = fw_utils.framework_version_from_tag(tag)
+        assert "1.5-1" == version
+
+
+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_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"