fix: Return ARM XGB/SKLearn tags if image_scope is inference_graviton

Author: Mark Bunday

src/sagemaker/image_uris.py

@@ -34,6 +34,7 @@
 XGBOOST_FRAMEWORK = "xgboost"
 SKLEARN_FRAMEWORK = "sklearn"
 TRAINIUM_ALLOWED_FRAMEWORKS = "pytorch"
+INFERENCE_GRAVITON = "inference_graviton"
 
 
 @override_pipeline_parameter_var
@@ -146,8 +147,9 @@ def retrieve(
         )
 
     if training_compiler_config and (framework == HUGGING_FACE_FRAMEWORK):
+        final_image_scope = image_scope
         config = _config_for_framework_and_scope(
-            framework + "-training-compiler", image_scope, accelerator_type
+            framework + "-training-compiler", final_image_scope, accelerator_type
         )
     else:
         _framework = framework
@@ -234,6 +236,7 @@ def retrieve(
     tag = _get_image_tag(
         container_version,
         distribution,
+        final_image_scope,
         framework,
         inference_tool,
         instance_type,
@@ -266,6 +269,7 @@ def _get_instance_type_family(instance_type):
 def _get_image_tag(
     container_version,
     distribution,
+    final_image_scope,
     framework,
     inference_tool,
     instance_type,
@@ -276,9 +280,9 @@ def _get_image_tag(
 ):
     """Return image tag based on framework, container, and compute configuration(s)."""
     instance_type_family = _get_instance_type_family(instance_type)
-    if (
-        framework in (XGBOOST_FRAMEWORK, SKLEARN_FRAMEWORK)
-        and instance_type_family in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
+    if framework in (XGBOOST_FRAMEWORK, SKLEARN_FRAMEWORK) and (
+        instance_type_family in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
+        or final_image_scope == INFERENCE_GRAVITON
     ):
         version_to_arm64_tag_mapping = {
             "xgboost": {
@@ -375,7 +379,7 @@ def _get_final_image_scope(framework, instance_type, image_scope):
         framework in GRAVITON_ALLOWED_FRAMEWORKS
         and _get_instance_type_family(instance_type) in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
     ):
-        return "inference_graviton"
+        return INFERENCE_GRAVITON
     if image_scope is None and framework in (XGBOOST_FRAMEWORK, SKLEARN_FRAMEWORK):
         # Preserves backwards compatibility with XGB/SKLearn configs which no
         # longer define top-level "scope" keys after introducing support for

tests/unit/sagemaker/image_uris/test_graviton.py

@@ -89,7 +89,7 @@ def test_graviton_pytorch(graviton_pytorch_version):
     _test_graviton_framework_uris("pytorch", graviton_pytorch_version)
 
 
-def test_graviton_xgboost(graviton_xgboost_versions):
+def test_graviton_xgboost_instance_type_specified(graviton_xgboost_versions):
     for xgboost_version in graviton_xgboost_versions:
         for instance_type in GRAVITON_INSTANCE_TYPES:
             uri = image_uris.retrieve(
@@ -102,6 +102,19 @@ def test_graviton_xgboost(graviton_xgboost_versions):
             assert expected == uri
 
 
+def test_graviton_xgboost_image_scope_specified(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, image_scope="inference_graviton"
+            )
+            expected = (
+                "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:"
+                f"{xgboost_version}-arm64"
+            )
+            assert expected == uri
+
+
 def test_graviton_xgboost_unsupported_version(graviton_xgboost_unsupported_versions):
     for xgboost_version in graviton_xgboost_unsupported_versions:
         for instance_type in GRAVITON_INSTANCE_TYPES:
@@ -112,7 +125,7 @@ def test_graviton_xgboost_unsupported_version(graviton_xgboost_unsupported_versi
             assert f"Unsupported xgboost version: {xgboost_version}." in str(error)
 
 
-def test_graviton_sklearn(graviton_sklearn_versions):
+def test_graviton_sklearn_instance_type_specified(graviton_sklearn_versions):
     for sklearn_version in graviton_sklearn_versions:
         for instance_type in GRAVITON_INSTANCE_TYPES:
             uri = image_uris.retrieve(
@@ -125,6 +138,19 @@ def test_graviton_sklearn(graviton_sklearn_versions):
             assert expected == uri
 
 
+def test_graviton_sklearn_image_scope_specified(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, image_scope="inference_graviton"
+            )
+            expected = (
+                "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:"
+                f"{sklearn_version}-arm64-cpu-py3"
+            )
+            assert expected == uri
+
+
 def test_graviton_sklearn_unsupported_version(graviton_sklearn_unsupported_versions):
     for sklearn_version in graviton_sklearn_unsupported_versions:
         for instance_type in GRAVITON_INSTANCE_TYPES: