fix: fix xgboost image incorrect latest version warning (#1434)

Author: balajitummala

src/sagemaker/amazon/amazon_estimator.py

@@ -27,7 +27,12 @@
 from sagemaker.model import NEO_IMAGE_ACCOUNT
 from sagemaker.session import s3_input
 from sagemaker.utils import sagemaker_timestamp, get_ecr_image_uri_prefix
-from sagemaker.xgboost.defaults import XGBOOST_VERSION_1, XGBOOST_SUPPORTED_VERSIONS
+from sagemaker.xgboost.defaults import (
+    XGBOOST_LATEST_VERSION,
+    XGBOOST_SUPPORTED_VERSIONS,
+    XGBOOST_VERSION_1,
+    XGBOOST_VERSION_EQUIVALENTS,
+)
 from sagemaker.xgboost.estimator import get_xgboost_image_uri
 
 logger = logging.getLogger(__name__)
@@ -611,24 +616,46 @@ def get_image_uri(region_name, repo_name, repo_version=1):
         repo_version:
     """
     if repo_name == "xgboost":
+        if not _is_latest_xgboost_version(repo_version):
+            logging.warning(
+                "There is a more up to date SageMaker XGBoost image. "
+                "To use the newer image, please set 'repo_version'="
+                "'%s'. For example:\n"
+                "\tget_image_uri(region, 'xgboost', '%s').",
+                XGBOOST_LATEST_VERSION,
+                XGBOOST_LATEST_VERSION,
+            )
+
         if repo_version in ["0.90", "0.90-1", "0.90-1-cpu-py3"]:
             return get_xgboost_image_uri(region_name, XGBOOST_VERSION_1)
 
         supported_version = [
             version
             for version in XGBOOST_SUPPORTED_VERSIONS
-            if repo_version in (version, version + "-cpu-py3")
+            if repo_version in _generate_version_equivalents(version)
         ]
         if supported_version:
             return get_xgboost_image_uri(region_name, supported_version[0])
 
-        logging.warning(
-            "There is a more up to date SageMaker XGBoost image. "
-            "To use the newer image, please set 'repo_version'="
-            "'%s'. For example:\n"
-            "\tget_image_uri(region, 'xgboost', '%s').",
-            XGBOOST_VERSION_1,
-            XGBOOST_VERSION_1,
-        )
     repo = "{}:{}".format(repo_name, repo_version)
     return "{}/{}".format(registry(region_name, repo_name), repo)
+
+
+def _is_latest_xgboost_version(repo_version):
+    """Compare xgboost image version with latest version
+
+    Args:
+        repo_version:
+    """
+    if repo_version in (1, "latest"):
+        return False
+    return repo_version in _generate_version_equivalents(XGBOOST_LATEST_VERSION)
+
+
+def _generate_version_equivalents(version):
+    """Returns a list of version equivalents for XGBoost
+
+    Args:
+        version:
+    """
+    return [version + suffix for suffix in XGBOOST_VERSION_EQUIVALENTS] + [version]

src/sagemaker/xgboost/defaults.py

@@ -17,3 +17,4 @@
 XGBOOST_VERSION_1 = "0.90-1"
 XGBOOST_LATEST_VERSION = "0.90-2"
 XGBOOST_SUPPORTED_VERSIONS = [XGBOOST_VERSION_1, XGBOOST_LATEST_VERSION]
+XGBOOST_VERSION_EQUIVALENTS = ["-cpu-py3"]

tests/unit/test_amazon_estimator.py

@@ -24,7 +24,9 @@
     registry,
     get_image_uri,
     FileSystemRecordSet,
+    _is_latest_xgboost_version,
 )
+from sagemaker.xgboost.defaults import XGBOOST_LATEST_VERSION, XGBOOST_SUPPORTED_VERSIONS
 
 COMMON_ARGS = {"role": "myrole", "train_instance_count": 1, "train_instance_type": "ml.c4.xlarge"}
 
@@ -474,3 +476,13 @@ def test_regitry_throws_error_if_mapping_does_not_exist_for_default_algorithm():
     with pytest.raises(ValueError) as error:
         registry("broken_region_name")
     assert "Algorithm (None) is unsupported for region (broken_region_name)." in str(error)
+
+
+def test_is_latest_xgboost_version():
+    for version in XGBOOST_SUPPORTED_VERSIONS:
+        if version != XGBOOST_LATEST_VERSION:
+            assert _is_latest_xgboost_version(version) is False
+
+    assert _is_latest_xgboost_version("0.90-1-cpu-py3") is False
+
+    assert _is_latest_xgboost_version(XGBOOST_LATEST_VERSION) is True