change: Allow either instance_type or instance_group to be defined in distributed training config.

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)}."
-                )
+                if self.instance_groups is not None:
+                    for instance in self.instance_groups:
+                        if "p5.48xlarge" in instance._to_request_dict()["InstanceType"]:
+                            p5_enabled = True
+
             img_uri = "" if self.image_uri is None else self.image_uri
             for unsupported_image in Framework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
                 if (