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,15 @@ 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-form: "315553699071.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1-cpu-py3"
+    # 2. Short-form: "315553699071.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1"
+    # Both tags point to the same image and the image postfixed with cpu can also supports gpu
+    short_xgboost_name_pattern = re.compile(r"^sagemaker-(xgboost):(.*)$")
+    short_xgboost_match = short_xgboost_name_pattern.match(sagemaker_match.group(9))
+    if short_xgboost_match is not None:
+        return (short_xgboost_match.group(1), "py3", short_xgboost_match.group(2), None)
     return None, None, None, None