change: don't require instance_type for image_uris.retrieve() if only one option

Author: laurenyu

src/sagemaker/image_uris.py

@@ -173,6 +173,10 @@ def _processor(instance_type, available_processors):
         logger.info("Ignoring unnecessary instance type: %s.", instance_type)
         return None
 
+    if len(available_processors) == 1 and not instance_type:
+        logger.info("Defaulting to only supported image scope: %s.", available_processors[0])
+        return available_processors[0]
+
     if not instance_type:
         raise ValueError(
             "Empty SageMaker instance type. For options, see: "

tests/unit/sagemaker/image_uris/test_retrieve.py

@@ -522,6 +522,22 @@ def test_retrieve_processor_type_from_version_specific_processor_config(config_f
     assert "123412341234.dkr.ecr.us-west-2.amazonaws.com/dummy:1.1.0-py3" == uri
 
 
+@patch("sagemaker.image_uris.config_for_framework")
+def test_retrieve_default_processor_type_if_possible(config_for_framework):
+    config = copy.deepcopy(BASE_CONFIG)
+    config["processors"] = ["cpu"]
+    config_for_framework.return_value = config
+
+    uri = image_uris.retrieve(
+        framework="useless-string",
+        version="1.0.0",
+        py_version="py3",
+        region="us-west-2",
+        image_scope="training",
+    )
+    assert "123412341234.dkr.ecr.us-west-2.amazonaws.com/dummy:1.0.0-cpu-py3" == uri
+
+
 @patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
 def test_retrieve_unsupported_processor_type(config_for_framework):
     with pytest.raises(ValueError) as e:

tests/unit/sagemaker/image_uris/test_xgboost.py

@@ -72,11 +72,7 @@
 def test_xgboost_framework(xgboost_framework_version):
     for region in regions.regions():
         uri = image_uris.retrieve(
-            framework="xgboost",
-            region=region,
-            version=xgboost_framework_version,
-            py_version="py3",
-            instance_type="ml.c4.xlarge",
+            framework="xgboost", region=region, version=xgboost_framework_version, py_version="py3",
         )
 
         expected = expected_uris.framework_uri(