Add support for xgboost version 0.90-2 (aws#258)

* change: add support for xgboost version 0.90-2

* fix: remove duplicate rule in test_debugger.py:test_mxnet_with_all_rules_and_configs

Author: aws-patlin

src/sagemaker/amazon/amazon_estimator.py

@@ -27,8 +27,8 @@
 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.estimator import get_xgboost_image_uri
-from sagemaker.xgboost.defaults import XGBOOST_LATEST_VERSION
 
 logger = logging.getLogger(__name__)
 
@@ -559,13 +559,23 @@ def get_image_uri(region_name, repo_name, repo_version=1):
     """
     if repo_name == "xgboost":
         if repo_version in ["0.90", "0.90-1", "0.90-1-cpu-py3"]:
-            return get_xgboost_image_uri(region_name, XGBOOST_LATEST_VERSION)
+            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 supported_version:
+            return get_xgboost_image_uri(region_name, supported_version[0])
+
         logging.warning(
-            "There is a more up to date SageMaker XGBoost image."
+            "There is a more up to date SageMaker XGBoost image. "
             "To use the newer image, please set 'repo_version'="
-            "'0.90-1. For example:\n"
+            "'%s'. For example:\n"
             "\tget_image_uri(region, 'xgboost', '%s').",
-            XGBOOST_LATEST_VERSION,
+            XGBOOST_VERSION_1,
+            XGBOOST_VERSION_1,
         )
     repo = "{}:{}".format(repo_name, repo_version)
     return "{}/{}".format(registry(region_name, repo_name), repo)

src/sagemaker/xgboost/defaults.py

@@ -14,5 +14,6 @@
 from __future__ import absolute_import
 
 XGBOOST_NAME = "xgboost"
-XGBOOST_LATEST_VERSION = "0.90-1"
-XGBOOST_SUPPORTED_VERSIONS = [XGBOOST_LATEST_VERSION]
+XGBOOST_VERSION_1 = "0.90-1"
+XGBOOST_LATEST_VERSION = "0.90-2"
+XGBOOST_SUPPORTED_VERSIONS = [XGBOOST_VERSION_1, XGBOOST_LATEST_VERSION]

tests/integ/test_debugger.py

@@ -345,7 +345,6 @@
 #         rules = [
 #             Rule.sagemaker(rule_configs.vanishing_gradient()),
 #             Rule.sagemaker(rule_configs.all_zero()),
-#             Rule.sagemaker(rule_configs.check_input_images()),
 #             Rule.sagemaker(rule_configs.similar_across_runs()),
 #             Rule.sagemaker(rule_configs.weight_update_ratio()),
 #             Rule.sagemaker(rule_configs.exploding_tensor()),

tests/unit/test_amazon_estimator.py

@@ -423,3 +423,9 @@ def test_get_xgboost_image_uri():
         updated_xgb_image_uri
         == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:0.90-1-cpu-py3"
     )
+
+    updated_xgb_image_uri_v2 = get_image_uri(REGION, "xgboost", "0.90-2")
+    assert (
+        updated_xgb_image_uri_v2
+        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:0.90-2-cpu-py3"
+    )

tests/unit/test_xgboost.py

@@ -175,7 +175,7 @@ def test_create_model(sagemaker_session):
         entry_point=SCRIPT_PATH,
         framework_version=XGBOOST_LATEST_VERSION,
     )
-    default_image_uri = _get_full_cpu_image_uri("0.90-1")
+    default_image_uri = _get_full_cpu_image_uri(XGBOOST_LATEST_VERSION)
     model_values = xgboost_model.prepare_container_def(CPU)
     assert model_values["Image"] == default_image_uri