fix: Fixing the logic to return instanceCount for heterogeneousClusters (aws#3363)

We recently started supporting feature for heterogeneousClusters and missed
this part where we return instanceCount from ResourceConfig -> InstanceType
For heterogeneousCluster jobs the total instanceCount is sum total of the
instanceCount of all the instance groups. Fixing the same.

Author: asumitamazon

src/sagemaker/session.py

@@ -4987,7 +4987,12 @@ def _rule_statuses_changed(current_statuses, last_statuses):
 def _logs_init(sagemaker_session, description, job):
     """Placeholder docstring"""
     if job == "Training":
-        instance_count = description["ResourceConfig"]["InstanceCount"]
+        if description["ResourceConfig"]["InstanceCount"] is not None:
+            instance_count = description["ResourceConfig"]["InstanceCount"]
+        else:
+            instance_count = 0
+            for instanceGroup in description["ResourceConfig"]["InstanceGroups"]:
+                instance_count += instanceGroup["InstanceCount"]
     elif job == "Transform":
         instance_count = description["TransformResources"]["InstanceCount"]
     elif job == "Processing":