feat: auto-select container version for p4d and smdistributed (#517)

Author: ajaykarpur

src/sagemaker/estimator.py

@@ -2297,13 +2297,18 @@ def training_image_uri(self):
         """
         if self.image_uri:
             return self.image_uri
+        if hasattr(self, "distribution"):
+            distribution = self.distribution  # pylint: disable=no-member
+        else:
+            distribution = None
         return image_uris.retrieve(
             self._framework_name,
             self.sagemaker_session.boto_region_name,
             instance_type=self.instance_type,
             version=self.framework_version,  # pylint: disable=no-member
             py_version=self.py_version,  # pylint: disable=no-member
             image_scope="training",
+            distribution=distribution,
         )
 
     @classmethod

src/sagemaker/image_uris.py

@@ -35,6 +35,7 @@ def retrieve(
     accelerator_type=None,
     image_scope=None,
     container_version=None,
+    distribution=None,
 ):
     """Retrieves the ECR URI for the Docker image matching the given arguments.
 
@@ -54,6 +55,8 @@ def retrieve(
             Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
             ``image_scope`` is ignored.
         container_version (str): the version of docker image
+        distribution (dict): A dictionary with information on how to run distributed training
+            (default: None).
 
     Returns:
         str: the ECR URI for the corresponding SageMaker Docker image.
@@ -77,10 +80,25 @@ def retrieve(
     processor = _processor(
         instance_type, config.get("processors") or version_config.get("processors")
     )
+
     tag = _format_tag(
-        version_config.get("tag_prefix", version), processor, py_version, container_version
+        version_config.get("tag_prefix", version),
+        processor,
+        py_version,
+        container_version,
     )
 
+    if _should_auto_select_container_version(instance_type, distribution):
+        container_versions = {
+            "tensorflow-2.3-gpu-py37": "cu110-ubuntu18.04-v3",
+            "tensorflow-1.15-gpu-py37": "cu110-ubuntu18.04-v8",
+            "mxnet-1.8-gpu-py37": "cu110-ubuntu16.04-v1",
+            "pytorch-1.6-gpu-py36": "cu110-ubuntu18.04-v3",
+        }
+        key = "-".join([framework, tag])
+        if key in container_versions:
+            tag = "-".join([tag, container_versions[key]])
+
     if tag:
         repo += ":{}".format(tag)
 
@@ -217,6 +235,23 @@ def _processor(instance_type, available_processors):
     return processor
 
 
+def _should_auto_select_container_version(instance_type, distribution):
+    """Returns a boolean that indicates whether to use an auto-selected container version."""
+    p4d = False
+    if instance_type:
+        # looks for either "ml.<family>.<size>" or "ml_<family>"
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match:
+            family = match[1]
+            p4d = family == "p4d"
+
+    smdistributed = False
+    if distribution:
+        smdistributed = "smdistributed" in distribution
+
+    return p4d or smdistributed
+
+
 def _validate_py_version_and_set_if_needed(py_version, version_config, framework):
     """Checks if the Python version is one of the supported versions."""
     if "repository" in version_config:

tests/unit/sagemaker/image_uris/test_retrieve.py

@@ -538,6 +538,21 @@ def test_retrieve_default_processor_type_if_possible(config_for_framework):
     assert "123412341234.dkr.ecr.us-west-2.amazonaws.com/dummy:1.0.0-cpu-py3" == uri
 
 
+def test_retrieve_auto_selected_container_version():
+    uri = image_uris.retrieve(
+        framework="tensorflow",
+        region="us-west-2",
+        version="2.3",
+        py_version="py37",
+        instance_type="ml.p4d.24xlarge",
+        image_scope="training",
+    )
+    assert (
+        "763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.3-gpu-py37-cu110-ubuntu18.04-v3"
+        == 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: