feature: heterogeneous cluster set up in distribution config

yl-to · jessicazhu3 · navinsoni · commit 1dc4cb0c4ebd · 2022-07-07T16:43:15.000-07:00
Co-authored-by: Jessica Zhu &lt;106775307+jessicazhu3@users.noreply.github.com&gt;
diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py
@@ -3111,6 +3111,13 @@ def _distribution_configuration(self, distribution):
         """
         distribution_config = {}
 
+        mpi_enabled = False
+        smdataparallel_enabled = False
+        if "instance_groups" in distribution:
+            distribution_config["sagemaker_distribution_instance_groups"] = distribution[
+                "instance_groups"
+            ]
+
         if "parameter_server" in distribution:
             ps_enabled = distribution.get("parameter_server").get("enabled", False)
             distribution_config[self.LAUNCH_PS_ENV_NAME] = ps_enabled
@@ -3146,6 +3153,13 @@ def _distribution_configuration(self, distribution):
                     "dataparallel"
                 ].get("custom_mpi_options", "")
 
+        if not (mpi_enabled or smdataparallel_enabled) and distribution_config.get(
+            "sagemaker_distribution_instance_groups"
+        ) not in [None, []]:
+            raise ValueError(
+                "Don't set training instance groups while no distribution strategies enabled!"
+            )
+
         return distribution_config
 
 
diff --git a/src/sagemaker/fw_utils.py b/src/sagemaker/fw_utils.py
@@ -27,7 +27,7 @@
 import sagemaker.utils
 from sagemaker.workflow import is_pipeline_variable
 
-from sagemaker.deprecations import renamed_warning
+from sagemaker.deprecations import renamed_warning, renamed_kwargs
 
 logger = logging.getLogger(__name__)
 
@@ -600,6 +600,106 @@ def _validate_smdataparallel_args(
         raise ValueError(err_msg)
 
 
+def validate_distribution(
+    distribution, instance_groups, framework_name, framework_version, py_version, image_uri, kwargs
+):
+    """Check if distribution strategy is correctly invoked by the user.
+
+    Currently, check for `dataparallel`, `modelparallel` and heterogeneous cluster set up.
+    Validate if the user requested strategy is supported.
+
+    Args:
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) For example:
+
+            .. code:: python
+
+                {
+                    "smdistributed": {
+                        "dataparallel": {
+                            "enabled": True
+                        }
+                    }
+                }
+        instance_groups ([InstanceGroup]): A list contains instance groups used for training.
+        framework_name (str): A string representing the name of framework selected.
+        framework_version (str): A string representing the framework version selected.
+        py_version (str): A string representing the python version selected.
+        image_uri (str): A string representing a Docker image URI.
+        kwargs(dict): Additional kwargs passed to this function
+
+    Returns:
+        distribution(dict): updated dictionary with validated information
+            to enable distributed training.
+
+    Raises:
+        ValueError: if distribution dictionary isn't correctly formatted or
+            multiple strategies are requested simultaneously or
+            an unsupported strategy is requested or
+            strategy-specific inputs are incorrect/unsupported or
+            heterogeneous cluster set up is incorrect
+    """
+    train_instance_groups = distribution.get("instance_groups", [])
+    if instance_groups is None:
+        if len(train_instance_groups) >= 1:
+            # if estimator's instance_groups is not defined but
+            # train_instance_groups are specified in distribution
+            raise ValueError("Instance groups not specified in the estimator !")
+    else:
+        if len(train_instance_groups) > len(instance_groups):
+            # if train_instance_groups in distribution are more than estimator's instance_groups
+            raise ValueError("Train instance groups oversubscribed !")
+        if len(instance_groups) == 1 and len(train_instance_groups) == 0:
+            # if just one instance_group but it is not specified in distribution, we set it for user
+            train_instance_groups = instance_groups
+        elif len(instance_groups) > 1 and len(train_instance_groups) != 1:
+            # currently we just support one train instance group
+            raise ValueError("Distribution should only contain one instance group name !")
+
+    if len(train_instance_groups) != 0:
+        # in this case, we are handling a heterogeneous cluster training job
+        instance_group_names = []
+        for train_instance_group in train_instance_groups:
+            # in future version we will support multiple train_instance_groups, so use loop here
+            if train_instance_group not in instance_groups:
+                # check if train instance groups belongs to what user defined in estimator set up
+                raise ValueError(
+                    f"Invalid training instance group {train_instance_group.instance_group_name} !"
+                )
+            instance_type = train_instance_group.instance_type
+            validate_smdistributed(
+                instance_type=instance_type,
+                framework_name=framework_name,
+                framework_version=framework_version,
+                py_version=py_version,
+                distribution=distribution,
+                image_uri=image_uri,
+            )
+            warn_if_parameter_server_with_multi_gpu(
+                training_instance_type=instance_type, distribution=distribution
+            )
+            # get instance group names
+            instance_group_names.append(train_instance_group.instance_group_name)
+        distribution["instance_groups"] = instance_group_names
+    else:
+        # in this case, we are handling a normal training job (without heterogeneous cluster)
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
+        validate_smdistributed(
+            instance_type=instance_type,
+            framework_name=framework_name,
+            framework_version=framework_version,
+            py_version=py_version,
+            distribution=distribution,
+            image_uri=image_uri,
+        )
+        warn_if_parameter_server_with_multi_gpu(
+            training_instance_type=instance_type, distribution=distribution
+        )
+    return distribution
+
+
 def python_deprecation_warning(framework, latest_supported_version):
     """Placeholder docstring"""
     return PYTHON_2_DEPRECATION_WARNING.format(
diff --git a/src/sagemaker/pytorch/estimator.py b/src/sagemaker/pytorch/estimator.py
@@ -17,15 +17,13 @@
 
 from packaging.version import Version
 
-from sagemaker.deprecations import renamed_kwargs
 from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.fw_utils import (
     framework_name_from_image,
     framework_version_from_tag,
     python_deprecation_warning,
     validate_version_or_image_args,
-    warn_if_parameter_server_with_multi_gpu,
-    validate_smdistributed,
+    validate_distribution,
 )
 from sagemaker.pytorch import defaults
 from sagemaker.pytorch.model import PyTorchModel
@@ -196,24 +194,6 @@ def __init__(
         self.framework_version = framework_version
         self.py_version = py_version
 
-        if distribution is not None:
-            instance_type = renamed_kwargs(
-                "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
-            )
-
-            validate_smdistributed(
-                instance_type=instance_type,
-                framework_name=self._framework_name,
-                framework_version=framework_version,
-                py_version=py_version,
-                distribution=distribution,
-                image_uri=image_uri,
-            )
-
-            warn_if_parameter_server_with_multi_gpu(
-                training_instance_type=instance_type, distribution=distribution
-            )
-
         if "enable_sagemaker_metrics" not in kwargs:
             # enable sagemaker metrics for PT v1.3 or greater:
             if self.framework_version and Version(self.framework_version) >= Version("1.3"):
@@ -222,6 +202,17 @@ def __init__(
         super(PyTorch, self).__init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
+        if distribution is not None:
+            distribution = validate_distribution(
+                distribution,
+                self.instance_groups,
+                self._framework_name,
+                framework_version,
+                py_version,
+                image_uri,
+                kwargs,
+            )
+
         self.distribution = distribution or {}
 
     def hyperparameters(self):
diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py
@@ -183,25 +183,22 @@ def __init__(
         self.py_version = py_version
         self.instance_type = instance_type
 
-        if distribution is not None:
-            fw.warn_if_parameter_server_with_multi_gpu(
-                training_instance_type=instance_type, distribution=distribution
-            )
-            fw.validate_smdistributed(
-                instance_type=instance_type,
-                framework_name=self._framework_name,
-                framework_version=framework_version,
-                py_version=py_version,
-                distribution=distribution,
-                image_uri=image_uri,
-            )
-
         if "enable_sagemaker_metrics" not in kwargs:
             # enable sagemaker metrics for TF v1.15 or greater:
             if framework_version and version.Version(framework_version) >= version.Version("1.15"):
                 kwargs["enable_sagemaker_metrics"] = True
 
         super(TensorFlow, self).__init__(image_uri=image_uri, **kwargs)
+        if distribution is not None:
+            distribution = fw.validate_distribution(
+                distribution,
+                self.instance_groups,
+                self._framework_name,
+                framework_version,
+                py_version,
+                image_uri,
+                kwargs,
+            )
         self.model_dir = model_dir
         self.distribution = distribution or {}
 
diff --git a/tests/unit/sagemaker/tensorflow/test_estimator.py b/tests/unit/sagemaker/tensorflow/test_estimator.py
@@ -22,6 +22,7 @@
 
 from sagemaker.estimator import _TrainingJob
 from sagemaker.tensorflow import TensorFlow
+from sagemaker.instance_group import InstanceGroup
 from tests.unit import DATA_DIR
 
 SCRIPT_FILE = "dummy_script.py"
@@ -538,3 +539,24 @@ def test_custom_image(sagemaker_session):
     custom_image = "tensorflow:latest"
     tf = _build_tf(sagemaker_session, image_uri=custom_image)
     assert custom_image == tf.training_image_uri()
+
+
+def test_tf_heterogeneous_cluster_distribution_config(
+    sagemaker_session, tensorflow_training_version, tensorflow_training_py_version
+):
+    if version.Version(tensorflow_training_version) < version.Version("2.0"):
+        pytest.skip("This test is for TF 2.0 and higher.")
+
+    training_group = InstanceGroup("train_group", "ml.c4.xlarge", 1)
+    expected_return = {"mpi": {"enabled": True}, "instance_groups": ["train_group"]}
+    tf = _build_tf(
+        sagemaker_session,
+        framework_version=tensorflow_training_version,
+        py_version=tensorflow_training_py_version,
+        instance_groups=[training_group],
+        distribution={
+            "mpi": {"enabled": True},
+            "instance_groups": [training_group],
+        },
+    )
+    assert tf.distribution == expected_return
diff --git a/tests/unit/test_fw_utils.py b/tests/unit/test_fw_utils.py
diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py
