feature: support heterogeneous cluster for training

Co-authored-by: Navin Soni <navinsoni89@gmail.com>

Author: qidewenwhen

src/sagemaker/algorithm.py

@@ -147,19 +147,19 @@ def __init__(
         self.algorithm_arn = algorithm_arn
         super(AlgorithmEstimator, self).__init__(
             role,
-            instance_count,
-            instance_type,
-            volume_size,
-            volume_kms_key,
-            max_run,
-            input_mode,
-            output_path,
-            output_kms_key,
-            base_job_name,
-            sagemaker_session,
-            tags,
-            subnets,
-            security_group_ids,
+            instance_count=instance_count,
+            instance_type=instance_type,
+            volume_size=volume_size,
+            volume_kms_key=volume_kms_key,
+            max_run=max_run,
+            input_mode=input_mode,
+            output_path=output_path,
+            output_kms_key=output_kms_key,
+            base_job_name=base_job_name,
+            sagemaker_session=sagemaker_session,
+            tags=tags,
+            subnets=subnets,
+            security_group_ids=security_group_ids,
             model_uri=model_uri,
             model_channel_name=model_channel_name,
             metric_definitions=metric_definitions,

src/sagemaker/estimator.py

@@ -145,6 +145,7 @@ def __init__(
         code_location: Optional[str] = None,
         entry_point: Optional[str] = None,
         dependencies: Optional[List[Union[str]]] = None,
+        instance_groups=None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -156,9 +157,10 @@ def __init__(
                 artifacts. After the endpoint is created, the inference code
                 might use the IAM role, if it needs to access an AWS resource.
             instance_count (int): Number of Amazon EC2 instances to use
-                for training.
+                for training. Required if instance_groups is not set.
             instance_type (str): Type of EC2 instance to use for training,
-                for example, 'ml.c4.xlarge'.
+                for example, 'ml.c4.xlarge'. Required if instance_groups is
+                not set.
             volume_size (int): Size in GB of the EBS volume to use for
                 storing input data during training (default: 30). Must be large
                 enough to store training data if File Mode is used (which is the
@@ -424,7 +426,10 @@ def __init__(
                     >>>     |------ virtual-env
 
                 This is not supported with "local code" in Local Mode.
-
+            instance_groups (list[InstanceGroup]): Optional. List of InstanceGroup
+                for specifying different instance groups for heterogeneous cluster.
+                For example: [sagemaker.InstanceGroup('worker','ml.p3dn.24xlarge',64),
+                sagemaker.InstanceGroup('server','ml.c5n.18xlarge',64)]
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -442,12 +447,10 @@ def __init__(
             "train_volume_kms_key", "volume_kms_key", volume_kms_key, kwargs
         )
 
-        if instance_count is None or instance_type is None:
-            raise ValueError("Both instance_count and instance_type are required.")
-
         self.role = role
         self.instance_count = instance_count
         self.instance_type = instance_type
+        self.instance_groups = instance_groups
         self.volume_size = volume_size
         self.volume_kms_key = volume_kms_key
         self.max_run = max_run
@@ -2103,6 +2106,7 @@ def __init__(
         code_location: Optional[str] = None,
         entry_point: Optional[str] = None,
         dependencies: Optional[List[str]] = None,
+        instance_groups=None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -2115,9 +2119,10 @@ def __init__(
                 artifacts. After the endpoint is created, the inference code
                 might use the IAM role, if it needs to access an AWS resource.
             instance_count (int): Number of Amazon EC2 instances to use
-                for training.
+                for training. Required if instance_groups is not set.
             instance_type (str): Type of EC2 instance to use for training,
-                for example, 'ml.c4.xlarge'.
+                for example, 'ml.c4.xlarge'. Required if instance_groups is
+                not set.
             volume_size (int): Size in GB of the EBS volume to use for
                 storing input data during training (default: 30). Must be large
                 enough to store training data if File Mode is used (which is the
@@ -2379,13 +2384,18 @@ def __init__(
                     >>>     |------ virtual-env
 
                 This is not supported with "local code" in Local Mode.
+            instance_groups (list[InstanceGroup]): Optional. List of InstanceGroup
+                for specifying different instance groups for heterogeneous cluster.
+                For example: [sagemaker.InstanceGroup('worker','ml.p3dn.24xlarge',64),
+                sagemaker.InstanceGroup('server','ml.c5n.18xlarge',64)]
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
         super(Estimator, self).__init__(
             role,
             instance_count,
             instance_type,
+            instance_groups,
             volume_size,
             volume_kms_key,
             max_run,

src/sagemaker/inputs.py

@@ -35,6 +35,7 @@ def __init__(
         content_type=None,
         record_wrapping=None,
         s3_data_type="S3Prefix",
+        instance_groups=None,
         input_mode=None,
         attribute_names=None,
         target_attribute_name=None,
@@ -60,6 +61,8 @@ def __init__(
                 listing the S3 data to train on. Both the ManifestFile and
                 AugmentedManifestFile formats are described in the SageMaker API documentation:
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_S3DataSource.html
+            instance_groups (list[str]): Optional. List of InstanceGroupNames to send data to
+                (default: None). By default, data will be sent to all groups.
             input_mode (str): Optional override for this channel's input mode (default: None).
                 By default, channels will use the input mode defined on
                 ``sagemaker.estimator.EstimatorBase.input_mode``, but they will ignore
@@ -97,6 +100,8 @@ def __init__(
             self.config["ContentType"] = content_type
         if record_wrapping is not None:
             self.config["RecordWrapperType"] = record_wrapping
+        if instance_groups is not None:
+            self.config["DataSource"]["S3DataSource"]["InstanceGroupNames"] = instance_groups
         if input_mode is not None:
             self.config["InputMode"] = input_mode
         if attribute_names is not None:

src/sagemaker/instance_group.py

@@ -0,0 +1,50 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This file defines instance group for heterogeneous cluster."""
+from __future__ import absolute_import
+
+
+class InstanceGroup(object):
+    """Accepts instance group parameters for conversion to request dict.
+
+    The `_to_request_dict` provides a method to turn the parameters into a dict.
+    """
+
+    def __init__(
+        self,
+        instance_group_name=None,
+        instance_type=None,
+        instance_count=None,
+    ):
+        """Initialize a ``InstanceGroup`` instance.
+
+        InstanceGroup accepts instance group parameters and provides a method to turn
+        these parameters into a dictionary.
+
+        Args:
+            instance_group_name (str): Name of the instance group.
+            instance_type (str): Type of EC2 instance to use in the instance group,
+                for example, 'ml.c4.xlarge'.
+            instance_count (int): Number of EC2 instances to use in the instance group.
+        """
+        self.instance_group_name = instance_group_name
+        self.instance_type = instance_type
+        self.instance_count = instance_count
+
+    def _to_request_dict(self):
+        """Generates a request dictionary using the parameters provided to the class."""
+        return {
+            "InstanceGroupName": self.instance_group_name,
+            "InstanceType": self.instance_type,
+            "InstanceCount": self.instance_count,
+        }

src/sagemaker/job.py

@@ -74,6 +74,7 @@ def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
         resource_config = _Job._prepare_resource_config(
             estimator.instance_count,
             estimator.instance_type,
+            estimator.instance_groups,
             estimator.volume_size,
             estimator.volume_kms_key,
         )
@@ -283,15 +284,31 @@ def _prepare_output_config(s3_path, kms_key_id):
         return config
 
     @staticmethod
-    def _prepare_resource_config(instance_count, instance_type, volume_size, volume_kms_key):
+    def _prepare_resource_config(
+        instance_count, instance_type, instance_groups, volume_size, volume_kms_key
+    ):
         """Placeholder docstring"""
         resource_config = {
-            "InstanceCount": instance_count,
-            "InstanceType": instance_type,
             "VolumeSizeInGB": volume_size,
         }
         if volume_kms_key is not None:
             resource_config["VolumeKmsKeyId"] = volume_kms_key
+        if instance_groups is not None:
+            if instance_count is not None or instance_type is not None:
+                raise ValueError(
+                    "instance_count and instance_type cannot be set when instance_groups is set"
+                )
+
+            resource_config["InstanceGroups"] = [
+                group._to_request_dict() for group in instance_groups
+            ]
+        else:
+            if instance_count is None or instance_type is None:
+                raise ValueError(
+                    "instance_count and instance_type must be set if instance_groups is not set"
+                )
+            resource_config["InstanceCount"] = instance_count
+            resource_config["InstanceType"] = instance_type
 
         return resource_config
 

tests/unit/test_estimator.py

@@ -43,6 +43,7 @@
 from sagemaker.estimator import Estimator, EstimatorBase, Framework, _TrainingJob
 from sagemaker.fw_utils import PROFILER_UNSUPPORTED_REGIONS
 from sagemaker.inputs import ShuffleConfig
+from sagemaker.instance_group import InstanceGroup
 from sagemaker.model import FrameworkModel
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.predictor import Predictor
@@ -323,6 +324,31 @@ def test_framework_all_init_args(sagemaker_session):
     }
 
 
+def test_framework_with_heterogeneous_cluster(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args["resource_config"]["InstanceGroups"][0] == {
+        "InstanceGroupName": "group1",
+        "InstanceCount": 1,
+        "InstanceType": "ml.c4.xlarge",
+    }
+    assert args["resource_config"]["InstanceGroups"][1] == {
+        "InstanceGroupName": "group2",
+        "InstanceCount": 2,
+        "InstanceType": "ml.m4.xlarge",
+    }
+
+
 def test_framework_with_debugger_and_built_in_rule(sagemaker_session):
     debugger_built_in_rule_with_custom_args = Rule.sagemaker(
         base_config=rule_configs.stalled_training_rule(),

tests/unit/test_inputs.py

@@ -67,6 +67,44 @@ def test_training_input_all_arguments():
     assert result.config == expected
 
 
+def test_training_input_all_arguments_heterogeneous_cluster():
+    prefix = "pre"
+    distribution = "FullyReplicated"
+    compression = "Gzip"
+    content_type = "text/csv"
+    record_wrapping = "RecordIO"
+    s3_data_type = "Manifestfile"
+    instance_groups = ["data-server"]
+    input_mode = "Pipe"
+    result = TrainingInput(
+        s3_data=prefix,
+        distribution=distribution,
+        compression=compression,
+        input_mode=input_mode,
+        content_type=content_type,
+        record_wrapping=record_wrapping,
+        s3_data_type=s3_data_type,
+        instance_groups=instance_groups,
+    )
+
+    expected = {
+        "DataSource": {
+            "S3DataSource": {
+                "S3DataDistributionType": distribution,
+                "S3DataType": s3_data_type,
+                "S3Uri": prefix,
+                "InstanceGroupNames": instance_groups,
+            }
+        },
+        "CompressionType": compression,
+        "ContentType": content_type,
+        "RecordWrapperType": record_wrapping,
+        "InputMode": input_mode,
+    }
+
+    assert result.config == expected
+
+
 def test_file_system_input_default_access_mode():
     file_system_id = "fs-0a48d2a1"
     file_system_type = "EFS"