fix: image_uri does not need to be specified with instance_groups

Author: jessicazhu3

src/sagemaker/estimator.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import os
+import re
 import uuid
 from abc import ABCMeta, abstractmethod
 from typing import Any, Dict, Union, Optional, List
@@ -1520,6 +1521,39 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
                 init_params["max_wait"] = max_wait
         return init_params
 
+    def _get_instance_type(self):
+        """Determine the instance type to be used in the training_image_uri function.
+
+        Returns:
+            instance_type: The instance_type to be used.
+        """
+        if self.instance_type is not None:
+            return self.instance_type
+
+        if not isinstance(self.instance_groups, list) or len(self.instance_groups) == 0:
+            raise ValueError(
+                "instance_groups must be set if instance_type is not set and instance_groups "
+                "must be a list."
+            )
+
+        for instance_group in self.instance_groups:
+            instance_type = instance_group.instance_type
+            match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+
+            if match:
+                family = match[1]
+                if family[0] in ("g", "p"):
+                    return instance_type
+            else:
+                raise ValueError(
+                    "Invalid SageMaker instance type for training with heterogeneous clusters: {}. "
+                    "For options see: https://aws.amazon.com/sagemaker/pricing/instance-types".format(
+                        instance_type
+                    )
+                )
+
+        return self.instance_groups[0].instance_type
+
     def transformer(
         self,
         instance_count,
@@ -2903,7 +2937,7 @@ def training_image_uri(self, region=None):
             compiler_config=getattr(self, "compiler_config", None),
             tensorflow_version=getattr(self, "tensorflow_version", None),
             pytorch_version=getattr(self, "pytorch_version", None),
-            instance_type=self.instance_type,
+            instance_type=self._get_instance_type(),
         )
 
     @classmethod

tests/unit/test_estimator.py

@@ -1334,6 +1334,68 @@ def test_invalid_custom_code_bucket(sagemaker_session):
     assert "Expecting 's3' scheme" in str(error)
 
 
+def test_get_instance_type_gpu(sagemaker_session):
+    estimator = Estimator(
+        image_uri="some-image",
+        role="some_image",
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.p3.16xlarge", 2),
+        ],
+        sagemaker_session=sagemaker_session,
+        base_job_name="base_job_name",
+    )
+
+    assert "ml.p3.16xlarge" == estimator._get_instance_type()
+
+
+def test_get_instance_type_cpu(sagemaker_session):
+    estimator = Estimator(
+        image_uri="some-image",
+        role="some_image",
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.c5.xlarge", 2),
+        ],
+        sagemaker_session=sagemaker_session,
+        base_job_name="base_job_name",
+    )
+
+    assert "ml.c4.xlarge" == estimator._get_instance_type()
+
+
+def test_get_instance_type_no_instance_groups(sagemaker_session):
+    estimator = Estimator(
+        image_uri="some-image",
+        role="some_image",
+        instance_type="ml.c4.xlarge",
+        instance_count=1,
+        sagemaker_session=sagemaker_session,
+        base_job_name="base_job_name",
+    )
+
+    assert "ml.c4.xlarge" == estimator._get_instance_type()
+
+
+def test_get_instance_type_no_instance_groups_or_instance_type(sagemaker_session):
+    estimator = Estimator(
+        image_uri="some-image",
+        role="some_image",
+        instance_type=None,
+        instance_count=None,
+        instance_groups=None,
+        sagemaker_session=sagemaker_session,
+        base_job_name="base_job_name",
+    )
+    with pytest.raises(ValueError) as error:
+        estimator._get_instance_type()
+
+    assert (
+        "instance_groups must be set if instance_type is not set and instance_groups must be a list."
+        in str(error)
+    )
+
+
 def test_augmented_manifest(sagemaker_session):
     fw = DummyFramework(
         entry_point=SCRIPT_PATH,