breaking: use image_uris.retrieve() for XGBoost URIs

This also deprecates sagemaker.fw_utils.get_unsupported_framework_version_error(),
as well as many of the constant defined in sagemaker.xgboost.defaults.

Author: laurenyu

src/sagemaker/amazon/amazon_estimator.py

@@ -26,14 +26,6 @@
 from sagemaker.inputs import FileSystemInput, TrainingInput
 from sagemaker.model import NEO_IMAGE_ACCOUNT
 from sagemaker.utils import sagemaker_timestamp, get_ecr_image_uri_prefix
-from sagemaker.xgboost.defaults import (
-    XGBOOST_1P_VERSIONS,
-    XGBOOST_LATEST_VERSION,
-    XGBOOST_NAME,
-    XGBOOST_SUPPORTED_VERSIONS,
-    XGBOOST_VERSION_EQUIVALENTS,
-)
-from sagemaker.xgboost.estimator import get_xgboost_image_uri
 
 logger = logging.getLogger(__name__)
 
@@ -622,76 +614,5 @@ def get_image_uri(region_name, repo_name, repo_version=1):
         "in SageMaker Python SDK v2."
     )
 
-    repo_version = str(repo_version)
-
-    if repo_name == XGBOOST_NAME:
-
-        if repo_version in XGBOOST_1P_VERSIONS:
-            _warn_newer_xgboost_image()
-            return "{}/{}:{}".format(registry(region_name, repo_name), repo_name, repo_version)
-
-        if "-" not in repo_version:
-            xgboost_version_matches = [
-                version
-                for version in XGBOOST_SUPPORTED_VERSIONS
-                if repo_version == version.split("-")[0]
-            ]
-            if xgboost_version_matches:
-                # Assumes that XGBOOST_SUPPORTED_VERSION is sorted from oldest version to latest.
-                # When SageMaker version is not specified, we use the oldest one that matches
-                # XGBoost version for backward compatibility.
-                repo_version = xgboost_version_matches[0]
-
-        supported_framework_versions = [
-            version
-            for version in XGBOOST_SUPPORTED_VERSIONS
-            if repo_version in _generate_version_equivalents(version)
-        ]
-
-        if not supported_framework_versions:
-            raise ValueError(
-                "SageMaker XGBoost version {} is not supported. Supported versions: {}".format(
-                    repo_version, ", ".join(XGBOOST_SUPPORTED_VERSIONS)
-                )
-            )
-
-        if not _is_latest_xgboost_version(repo_version):
-            _warn_newer_xgboost_image()
-
-        return get_xgboost_image_uri(region_name, supported_framework_versions[-1])
-
     repo = "{}:{}".format(repo_name, repo_version)
     return "{}/{}".format(registry(region_name, repo_name), repo)
-
-
-def _warn_newer_xgboost_image():
-    """Print a warning when there is a newer XGBoost image"""
-    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, '%s', '%s').",
-        XGBOOST_LATEST_VERSION,
-        XGBOOST_NAME,
-        XGBOOST_LATEST_VERSION,
-    )
-
-
-def _is_latest_xgboost_version(repo_version):
-    """Compare xgboost image version with latest version
-
-    Args:
-        repo_version:
-    """
-    if repo_version in XGBOOST_1P_VERSIONS:
-        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/fw_utils.py

@@ -65,9 +65,6 @@
     "framework_version is required for script mode estimator. "
     "Please add framework_version={} to your constructor to avoid this error."
 )
-UNSUPPORTED_FRAMEWORK_VERSION_ERROR = (
-    "{} framework does not support version {}. Please use one of the following: {}."
-)
 
 VALID_PY_VERSIONS = ["py2", "py3", "py37"]
 VALID_EIA_FRAMEWORKS = [
@@ -637,25 +634,6 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
         logger.warning(PARAMETER_SERVER_MULTI_GPU_WARNING)
 
 
-def get_unsupported_framework_version_error(
-    framework_name, unsupported_version, supported_versions
-):
-    """Return error message for unsupported framework version.
-
-    This should also return the supported versions for customers.
-
-    :param framework_name:
-    :param unsupported_version:
-    :param supported_versions:
-    :return:
-    """
-    return UNSUPPORTED_FRAMEWORK_VERSION_ERROR.format(
-        framework_name,
-        unsupported_version,
-        ", ".join('"{}"'.format(version) for version in supported_versions),
-    )
-
-
 def python_deprecation_warning(framework, latest_supported_version):
     """
     Args:

src/sagemaker/xgboost/__init__.py

@@ -11,6 +11,6 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Placeholder docstring"""
-from sagemaker.xgboost.defaults import XGBOOST_NAME, XGBOOST_LATEST_VERSION  # noqa: F401
+from sagemaker.xgboost.defaults import XGBOOST_NAME  # noqa: F401
 from sagemaker.xgboost.estimator import XGBoost  # noqa: F401
 from sagemaker.xgboost.model import XGBoostModel, XGBoostPredictor  # noqa: F401

src/sagemaker/xgboost/defaults.py

@@ -14,20 +14,3 @@
 from __future__ import absolute_import
 
 XGBOOST_NAME = "xgboost"
-XGBOOST_1P_VERSIONS = ["1", "latest"]
-
-XGBOOST_VERSION_0_90_1 = "0.90-1"
-XGBOOST_VERSION_0_90_2 = "0.90-2"
-
-XGBOOST_LATEST_VERSION = "1.0-1"
-
-# XGBOOST_SUPPORTED_VERSIONS has XGBoost Framework versions sorted from oldest to latest
-XGBOOST_SUPPORTED_VERSIONS = [
-    XGBOOST_VERSION_0_90_1,
-    XGBOOST_VERSION_0_90_2,
-    XGBOOST_LATEST_VERSION,
-]
-
-# TODO: evaluate use of this constant. it's used in precisely one place in different a module
-# may possibly be unnecessary indirection
-XGBOOST_VERSION_EQUIVALENTS = ["-cpu-py3"]

src/sagemaker/xgboost/estimator.py

@@ -15,12 +15,11 @@
 
 import logging
 
+from sagemaker import image_uris
 from sagemaker.estimator import Framework, _TrainingJob
-from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import (
     framework_name_from_image,
     framework_version_from_tag,
-    get_unsupported_framework_version_error,
     UploadedCode,
 )
 from sagemaker.session import Session
@@ -31,12 +30,6 @@
 logger = logging.getLogger("sagemaker")
 
 
-def get_xgboost_image_uri(region, framework_version, py_version="py3"):
-    """Get XGBoost framework image URI"""
-    image_tag = "{}-{}-{}".format(framework_version, "cpu", py_version)
-    return default_framework_uri(XGBoost.__framework_name__, region, image_tag)
-
-
 class XGBoost(Framework):
     """Handle end-to-end training and deployment of XGBoost booster training or training using
     customer provided XGBoost entry point script."""
@@ -105,22 +98,17 @@ def __init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
 
-        if py_version == "py2":
-            raise AttributeError("XGBoost container does not support Python 2, please use Python 3")
         self.py_version = py_version
-
-        if framework_version in defaults.XGBOOST_SUPPORTED_VERSIONS:
-            self.framework_version = framework_version
-        else:
-            raise ValueError(
-                get_unsupported_framework_version_error(
-                    self.__framework_name__, framework_version, defaults.XGBOOST_SUPPORTED_VERSIONS
-                )
-            )
+        self.framework_version = framework_version
 
         if image_uri is None:
-            self.image_uri = get_xgboost_image_uri(
-                self.sagemaker_session.boto_region_name, framework_version
+            self.image_uri = image_uris.retrieve(
+                self.__framework_name__,
+                self.sagemaker_session.boto_region_name,
+                version=framework_version,
+                py_version=self.py_version,
+                instance_type=kwargs.get("instance_type"),
+                image_scope="training",
             )
 
     def create_model(

src/sagemaker/xgboost/model.py

@@ -16,9 +16,9 @@
 import logging
 
 import sagemaker
+from sagemaker import image_uris
 from sagemaker.deserializers import CSVDeserializer
 from sagemaker.fw_utils import model_code_key_prefix
-from sagemaker.fw_registry import default_framework_uri
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import NumpySerializer
@@ -100,9 +100,6 @@ def __init__(
             model_data, image_uri, role, entry_point, predictor_cls=predictor_cls, **kwargs
         )
 
-        if py_version == "py2":
-            raise AttributeError("XGBoost container does not support Python 2, please use Python 3")
-
         self.py_version = py_version
         self.framework_version = framework_version
         self.model_server_workers = model_server_workers
@@ -136,17 +133,21 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
             deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
         return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
 
-    def serving_image_uri(self, region_name, instance_type):  # pylint: disable=unused-argument
+    def serving_image_uri(self, region_name, instance_type):
         """Create a URI for the serving image.
 
         Args:
             region_name (str): AWS region where the image is uploaded.
-            instance_type (str): SageMaker instance type. This parameter is unused because
-                XGBoost supports only CPU.
+            instance_type (str): SageMaker instance type. Must be a CPU instance type.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
-
         """
-        image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
-        return default_framework_uri(self.__framework_name__, region_name, image_tag)
+        return image_uris.retrieve(
+            self.__framework_name__,
+            region_name,
+            version=self.framework_version,
+            py_version=self.py_version,
+            instance_type=instance_type,
+            image_scope="inference",
+        )

tests/conftest.py

@@ -285,16 +285,6 @@ def tf_full_py_version(tf_full_version):
     return "py37"
 
 
-@pytest.fixture(scope="module")
-def xgboost_full_version():
-    return "1.0-1"
-
-
-@pytest.fixture(scope="module")
-def xgboost_full_py_version():
-    return "py3"
-
-
 @pytest.fixture(scope="session")
 def cpu_instance_type(sagemaker_session, request):
     region = sagemaker_session.boto_session.region_name

tests/integ/test_airflow_config.py

@@ -549,13 +549,13 @@ def test_tf_airflow_config_uploads_data_source_to_s3(
 
 @pytest.mark.canary_quick
 def test_xgboost_airflow_config_uploads_data_source_to_s3(
-    sagemaker_session, cpu_instance_type, xgboost_full_version, xgboost_full_py_version
+    sagemaker_session, cpu_instance_type, xgboost_latest_version, xgboost_latest_py_version
 ):
     with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
         xgboost = XGBoost(
             entry_point=os.path.join(DATA_DIR, "dummy_script.py"),
-            framework_version=xgboost_full_version,
-            py_version=xgboost_full_py_version,
+            framework_version=xgboost_latest_version,
+            py_version=xgboost_latest_py_version,
             role=ROLE,
             sagemaker_session=sagemaker_session,
             instance_type=cpu_instance_type,

tests/unit/test_amazon_estimator.py

@@ -24,9 +24,7 @@
     registry,
     get_image_uri,
     FileSystemRecordSet,
-    _is_latest_xgboost_version,
 )
-from sagemaker.xgboost.defaults import XGBOOST_LATEST_VERSION, XGBOOST_SUPPORTED_VERSIONS
 
 COMMON_ARGS = {"role": "myrole", "instance_count": 1, "instance_type": "ml.c4.xlarge"}
 
@@ -449,73 +447,6 @@ def test_file_system_record_set_data_channel():
     assert actual == expected
 
 
-def test_get_xgboost_image_uri():
-    legacy_xgb_image_uri = get_image_uri(REGION, "xgboost")
-    assert legacy_xgb_image_uri == "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1"
-    legacy_xgb_image_uri = get_image_uri(REGION, "xgboost", 1)
-    assert legacy_xgb_image_uri == "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1"
-    legacy_xgb_image_uri = get_image_uri(REGION, "xgboost", "latest")
-    assert legacy_xgb_image_uri == "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:latest"
-
-    updated_xgb_image_uri = get_image_uri(REGION, "xgboost", "0.90-1")
-    assert (
-        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"
-    )
-
-    assert (
-        get_image_uri(REGION, "xgboost", "0.90-2-cpu-py3")
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:0.90-2-cpu-py3"
-    )
-    assert (
-        get_image_uri(REGION, "xgboost", "0.90")
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:0.90-1-cpu-py3"
-    )
-    assert (
-        get_image_uri(REGION, "xgboost", "1.0-1")
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.0-1-cpu-py3"
-    )
-    assert (
-        get_image_uri(REGION, "xgboost", "1.0-1-cpu-py3")
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.0-1-cpu-py3"
-    )
-    assert (
-        get_image_uri(REGION, "xgboost", "1.0")
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.0-1-cpu-py3"
-    )
-
-
-def test_get_xgboost_image_uri_warning_with_legacy(caplog):
-    get_image_uri(REGION, "xgboost", 1)
-    assert "There is a more up to date SageMaker XGBoost image." in caplog.text
-
-
-def test_get_xgboost_image_uri_warning_with_no_sagemaker_version(caplog):
-    get_image_uri(REGION, "xgboost", "0.90")
-    assert "There is a more up to date SageMaker XGBoost image." in caplog.text
-
-
-def test_get_xgboost_image_uri_no_warning_with_latest(caplog):
-    get_image_uri(REGION, "xgboost", XGBOOST_LATEST_VERSION.split("-")[0])
-    assert "There is a more up to date SageMaker XGBoost image." not in caplog.text
-
-
-def test_get_xgboost_image_uri_throws_error_for_unsupported_version():
-    with pytest.raises(ValueError) as error:
-        get_image_uri(REGION, "xgboost", "99.99-9")
-    assert "SageMaker XGBoost version 99.99-9 is not supported" in str(error)
-
-    with pytest.raises(ValueError) as error:
-        get_image_uri(REGION, "xgboost", "0.90-1-gpu-py3")
-    assert "SageMaker XGBoost version 0.90-1-gpu-py3 is not supported" in str(error)
-
-
 def test_regitry_throws_error_if_mapping_does_not_exist_for_lda():
     with pytest.raises(ValueError) as error:
         registry("cn-north-1", "lda")
@@ -528,14 +459,6 @@ def test_regitry_throws_error_if_mapping_does_not_exist_for_default_algorithm():
     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
-        else:
-            assert _is_latest_xgboost_version(version) is True
-
-
 def test_get_image_uri_warn(caplog):
     warning_message = (
         "'get_image_uri' method will be deprecated in favor of 'ImageURIProvider' class "