change: Allow either instance_type or instance_group to be defined in… (#4232)

Author: rohithn1

src/sagemaker/estimator.py

@@ -3818,6 +3818,7 @@ def _distribution_configuration(self, distribution):
 
         mpi_enabled = False
         smdataparallel_enabled = False
+        p5_enabled = False
         if "instance_groups" in distribution:
             distribution_config["sagemaker_distribution_instance_groups"] = distribution[
                 "instance_groups"
@@ -3862,10 +3863,11 @@ def _distribution_configuration(self, distribution):
             elif isinstance(self.instance_type, str):
                 p5_enabled = "p5.48xlarge" in self.instance_type
             else:
-                raise ValueError(
-                    "Invalid object type for instance_type argument. Expected "
-                    f"{type(str)} or {type(ParameterString)} but got {type(self.instance_type)}."
-                )
+                for instance in self.instance_groups:
+                    if "p5.48xlarge" in instance._to_request_dict().get("InstanceType", ()):
+                        p5_enabled = True
+                        break
+
             img_uri = "" if self.image_uri is None else self.image_uri
             for unsupported_image in Framework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
                 if (

tests/unit/test_estimator.py

@@ -402,6 +402,39 @@ def test_validate_smdistributed_backward_compat_p4_not_raises(sagemaker_session)
     f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED)
 
 
+def test_validate_smdistributed_instance_groups_raises(sagemaker_session):
+    instance_group_1 = InstanceGroup("train_group", "ml.p4d.24xlarge", 2)
+    instance_group_2 = InstanceGroup("train_group", "ml.p5.48xlarge", 2)
+    f = DummyFramework(
+        "some_script.py",
+        role="DummyRole",
+        instance_groups=[instance_group_1, instance_group_2],
+        sagemaker_session=sagemaker_session,
+        output_path="outputpath",
+        image_uri="some_acceptable_image",
+    )
+    # Testing instance_group with p5 raises exception
+    with pytest.raises(ValueError):
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED)
+    with pytest.raises(ValueError):
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED)
+
+
+def test_validate_smdistributed_instance_groups_not_raises(sagemaker_session):
+    instance_group_1 = InstanceGroup("train_group", "ml.p4d.24xlarge", 2)
+    f = DummyFramework(
+        "some_script.py",
+        role="DummyRole",
+        instance_groups=[instance_group_1],
+        sagemaker_session=sagemaker_session,
+        output_path="outputpath",
+        image_uri="some_acceptable_image",
+    )
+    # Testing instance_group without p5 does not raise exception
+    f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED)
+    f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED)
+
+
 def test_framework_all_init_args(sagemaker_session):
     f = DummyFramework(
         "my_script.py",