feature: Graviton support for XGB and SKLearn frameworks

Author: Mark Bunday

src/sagemaker/fw_utils.py

@@ -169,6 +169,43 @@ def validate_source_dir(script, directory):
 GRAVITON_ALLOWED_FRAMEWORKS = set(["tensorflow", "pytorch"])
 
 
+GRAVITON_ALLOWED_REGIONS = {
+    "af-south-1",
+    "ap-east-1",
+    "ap-northeast-1",
+    "ap-northeast-2",
+    "ap-northeast-3",
+    "ap-south-1",
+    "ap-southeast-1",
+    "ap-southeast-2",
+    "ap-southeast-3",
+    "ca-central-1",
+    "cn-north-1",
+    "cn-northwest-1",
+    "eu-central-1",
+    "eu-north-1",
+    "eu-south-1",
+    "eu-west-1",
+    "eu-west-2",
+    "eu-west-3",
+    "me-south-1",
+    "sa-east-1",
+    "us-east-1",
+    "us-east-2",
+    "us-gov-west-1",
+    "us-iso-east-1",
+    "us-west-1",
+    "us-west-2",
+}
+
+
+# TODO: Consider combining with GRAVITON_ALLOWED_FRAMEWORKS into a dictionary
+XGBOOST_GRAVITON_INFERENCE_ENABLED_VERSIONS = ["1.5-1", "1.3-1"]
+
+
+SKLEARN_GRAVITON_INFERENCE_ENABLED_VERSIONS = ["1.0-1"]
+
+
 def validate_source_code_input_against_pipeline_variables(
     entry_point: Optional[Union[str, PipelineVariable]] = None,
     source_dir: Optional[Union[str, PipelineVariable]] = None,

src/sagemaker/image_uris.py

@@ -25,12 +25,20 @@
 from sagemaker.jumpstart import artifacts
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.utilities import override_pipeline_parameter_var
-from sagemaker.fw_utils import GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY, GRAVITON_ALLOWED_FRAMEWORKS
+from sagemaker.fw_utils import (
+    GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY,
+    GRAVITON_ALLOWED_FRAMEWORKS,
+    GRAVITON_ALLOWED_REGIONS,
+    SKLEARN_GRAVITON_INFERENCE_ENABLED_VERSIONS,
+    XGBOOST_GRAVITON_INFERENCE_ENABLED_VERSIONS,
+)
 
 logger = logging.getLogger(__name__)
 
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
+XGBOOST_FRAMEWORK = "xgboost"
+SKLEARN_FRAMEWORK = "sklearn"
 
 
 @override_pipeline_parameter_var
@@ -244,6 +252,25 @@ def retrieve(
         if key in container_versions:
             tag = "-".join([tag, container_versions[key]])
 
+    if framework in (XGBOOST_FRAMEWORK, SKLEARN_FRAMEWORK):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1] in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY:
+            _validate_arg(region, GRAVITON_ALLOWED_REGIONS, "Graviton region")
+            if framework == XGBOOST_FRAMEWORK:
+                _validate_arg(
+                    version, 
+                    XGBOOST_GRAVITON_INFERENCE_ENABLED_VERSIONS, 
+                    "xgboost version for Graviton instances"
+                )
+                tag = f"{version}-arm64"
+            else:
+                _validate_arg(
+                    version, 
+                    SKLEARN_GRAVITON_INFERENCE_ENABLED_VERSIONS, 
+                    "sklearn version for Graviton instances"
+                )
+                tag = f"{version}-arm64-cpu-py3"
+
     if tag:
         repo += ":{}".format(tag)
 

tests/conftest.py

@@ -22,7 +22,7 @@
 from botocore.config import Config
 from packaging.version import Version
 
-from sagemaker import Session, image_uris, utils
+from sagemaker import Session, fw_utils, image_uris, utils
 from sagemaker.local import LocalSession
 from sagemaker.workflow.pipeline_context import PipelineSession, LocalPipelineSession
 
@@ -318,6 +318,16 @@ def graviton_pytorch_version():
     return "1.12.1"
 
 
+@pytest.fixture(scope="module")
+def graviton_xgboost_versions():
+    return fw_utils.XGBOOST_GRAVITON_INFERENCE_ENABLED_VERSIONS
+
+
+@pytest.fixture(scope="module")
+def graviton_sklearn_versions():
+    return fw_utils.SKLEARN_GRAVITON_INFERENCE_ENABLED_VERSIONS
+
+
 @pytest.fixture(scope="module")
 def huggingface_tensorflow_latest_training_py_version():
     return "py38"

tests/unit/sagemaker/image_uris/test_graviton.py

@@ -73,6 +73,31 @@ def test_graviton_pytorch(graviton_pytorch_version):
     _test_graviton_framework_uris("pytorch", graviton_pytorch_version)
 
 
+def test_graviton_xgboost(graviton_xgboost_versions):
+    for xgboost_version in graviton_xgboost_versions:
+        for instance_type in GRAVITON_INSTANCE_TYPES:
+            uri = image_uris.retrieve(
+                "xgboost", "us-west-2", version=xgboost_version, instance_type=instance_type
+            )
+            expected = (
+                "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:"
+                f"{xgboost_version}-arm64")
+            assert expected == uri
+
+
+def test_graviton_sklearn(graviton_sklearn_versions):
+    for sklearn_version in graviton_sklearn_versions:
+        for instance_type in GRAVITON_INSTANCE_TYPES:
+            uri = image_uris.retrieve(
+                "sklearn", "us-west-2", version=sklearn_version, instance_type=instance_type
+            )
+            expected = (
+                "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:"
+                f"{sklearn_version}-arm64-cpu-py3"
+            )
+            assert expected == uri
+
+
 def _expected_graviton_framework_uri(framework, version, region):
     return expected_uris.graviton_framework_uri(
         "{}-inference-graviton".format(framework),