feature: Add mpi support for mxnet estimator api (#1581)

Author: ChaiBapchya

src/sagemaker/mxnet/estimator.py

@@ -99,8 +99,44 @@ def __init__(
                     * ``custom-image:latest``
 
             distributions (dict): A dictionary with information on how to run distributed
-                training (default: None). To have parameter servers launched for training,
-                set this value to be ``{'parameter_server': {'enabled': True}}``.
+                training (default: None). Currently we support distributed training with
+                parameter server and MPI [Horovod].
+                To enable parameter server use the following setup:
+
+                .. code:: python
+
+                    {
+                        'parameter_server':
+                        {
+                            'enabled': True
+                        }
+                    }
+
+                To enable MPI:
+
+                .. code:: python
+
+                    {
+                        'mpi':
+                        {
+                            'enabled': True
+                        }
+                    }
+
+                Option parameters within ``mpi`` are ``processes_per_host``
+                and ``custom_mpi_options``.
+
+                .. code:: python
+
+                    {
+                        'mpi':
+                        {
+                            'enabled': True,
+                            'processes_per_host': 2,
+                            'custom_mpi_options': '-verbose --NCCL_DEBUG=INFO'
+                        }
+                    }
+
             **kwargs: Additional kwargs passed to the
                 :class:`~sagemaker.estimator.Framework` constructor.
 
@@ -159,6 +195,20 @@ def _configure_distribution(self, distributions):
             enabled = distributions["parameter_server"].get("enabled", False)
             self._hyperparameters[self.LAUNCH_PS_ENV_NAME] = enabled
 
+        if "mpi" in distributions:
+            mpi_dict = distributions["mpi"]
+            mpi_enabled = mpi_dict.get("enabled", False)
+            self._hyperparameters[self.LAUNCH_MPI_ENV_NAME] = mpi_enabled
+
+            if mpi_dict.get("processes_per_host"):
+                self._hyperparameters[self.MPI_NUM_PROCESSES_PER_HOST] = mpi_dict.get(
+                    "processes_per_host"
+                )
+
+                self._hyperparameters[self.MPI_CUSTOM_MPI_OPTIONS] = mpi_dict.get(
+                    "custom_mpi_options", ""
+                )
+
     def create_model(
         self,
         model_server_workers=None,

tests/conftest.py

@@ -287,6 +287,11 @@ def cpu_instance_type(sagemaker_session, request):
         return "ml.m4.xlarge"
 
 
+@pytest.fixture(scope="module")
+def gpu_instance_type(request):
+    return "ml.p2.xlarge"
+
+
 @pytest.fixture(scope="session")
 def inf_instance_type(sagemaker_session, request):
     return "ml.inf1.xlarge"

tests/data/horovod/hvd_mnist_mxnet.py

@@ -0,0 +1,14 @@
+import json
+import os
+
+import horovod.mxnet as hvd
+
+if __name__ == "__main__":
+
+    hvd.init()
+
+    with open(os.path.join("/opt/ml/model/rank-%s" % (hvd.rank())), "w+") as f:
+        basic_info = {"rank": hvd.rank(), "size": hvd.size()}
+
+        json.dump(basic_info, f)
+        print('Saved file "rank-%s": %s' % (hvd.rank(), basic_info))

tests/integ/test_horovod.py

@@ -15,10 +15,10 @@
 import json
 import os
 import tarfile
-from six.moves.urllib.parse import urlparse
 
 import boto3
 import pytest
+from six.moves.urllib.parse import urlparse
 
 import sagemaker.utils
 import tests.integ as integ
@@ -28,11 +28,6 @@
 horovod_dir = os.path.join(os.path.dirname(__file__), "..", "data", "horovod")
 
 
-@pytest.fixture(scope="module")
-def gpu_instance_type(request):
-    return "ml.p2.xlarge"
-
-
 @pytest.mark.canary_quick
 def test_hvd_cpu(sagemaker_session, cpu_instance_type, tmpdir):
     _create_and_fit_estimator(sagemaker_session, cpu_instance_type, tmpdir)
@@ -46,41 +41,6 @@ def test_hvd_gpu(sagemaker_session, gpu_instance_type, tmpdir):
     _create_and_fit_estimator(sagemaker_session, gpu_instance_type, tmpdir)
 
 
-@pytest.mark.local_mode
-@pytest.mark.parametrize("instances, processes", [[1, 2], (2, 1), (2, 2)])
-def test_horovod_local_mode(sagemaker_local_session, instances, processes, tmpdir):
-    output_path = "file://%s" % tmpdir
-    job_name = sagemaker.utils.unique_name_from_base("tf-horovod")
-    estimator = TensorFlow(
-        entry_point=os.path.join(horovod_dir, "hvd_basic.py"),
-        role="SageMakerRole",
-        train_instance_count=2,
-        train_instance_type="local",
-        sagemaker_session=sagemaker_local_session,
-        py_version=integ.PYTHON_VERSION,
-        script_mode=True,
-        output_path=output_path,
-        framework_version="1.12",
-        distributions={"mpi": {"enabled": True, "processes_per_host": processes}},
-    )
-
-    with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
-        estimator.fit(job_name=job_name)
-
-        tmp = str(tmpdir)
-        extract_files(output_path.replace("file://", ""), tmp)
-
-        size = instances * processes
-
-        for rank in range(size):
-            assert read_json("rank-%s" % rank, tmp)["rank"] == rank
-
-
-def extract_files(output_path, tmpdir):
-    with tarfile.open(os.path.join(output_path, "model.tar.gz")) as tar:
-        tar.extractall(tmpdir)
-
-
 def read_json(file, tmp):
     with open(os.path.join(tmp, file)) as f:
         return json.load(f)

tests/integ/test_horovod_mx.py

@@ -0,0 +1,80 @@
+# Copyright 2017-2020 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.
+from __future__ import absolute_import
+
+import json
+import os
+import tarfile
+
+import boto3
+import pytest
+from six.moves.urllib.parse import urlparse
+
+import sagemaker.utils
+import tests.integ as integ
+from sagemaker.mxnet import MXNet
+from tests.integ import timeout
+
+horovod_dir = os.path.join(os.path.dirname(__file__), "..", "data", "horovod")
+
+
+@pytest.mark.canary_quick
+def test_hvd_cpu(sagemaker_session, cpu_instance_type, tmpdir):
+    _create_and_fit_estimator(sagemaker_session, cpu_instance_type, tmpdir)
+
+
+@pytest.mark.canary_quick
+@pytest.mark.skipif(
+    integ.test_region() in integ.TRAINING_NO_P2_REGIONS, reason="no ml.p2 instances in this region"
+)
+def test_hvd_gpu(sagemaker_session, gpu_instance_type, tmpdir):
+    _create_and_fit_estimator(sagemaker_session, gpu_instance_type, tmpdir)
+
+
+def read_json(file, tmp):
+    with open(os.path.join(tmp, file)) as f:
+        return json.load(f)
+
+
+def extract_files_from_s3(s3_url, tmpdir, sagemaker_session):
+    parsed_url = urlparse(s3_url)
+    s3 = boto3.resource("s3", region_name=sagemaker_session.boto_region_name)
+
+    model = os.path.join(tmpdir, "model")
+    s3.Bucket(parsed_url.netloc).download_file(parsed_url.path.lstrip("/"), model)
+
+    with tarfile.open(model, "r") as tar_file:
+        tar_file.extractall(tmpdir)
+
+
+def _create_and_fit_estimator(sagemaker_session, instance_type, tmpdir):
+    job_name = sagemaker.utils.unique_name_from_base("mx-horovod")
+    estimator = MXNet(
+        entry_point=os.path.join(horovod_dir, "hvd_mnist_mxnet.py"),
+        role="SageMakerRole",
+        train_instance_count=2,
+        train_instance_type=instance_type,
+        sagemaker_session=sagemaker_session,
+        py_version=integ.PYTHON_VERSION,
+        framework_version="1.6.0",
+        distributions={"mpi": {"enabled": True}},
+    )
+
+    with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        estimator.fit(job_name=job_name)
+
+        tmp = str(tmpdir)
+        extract_files_from_s3(estimator.model_data, tmp, sagemaker_session)
+
+        for rank in range(2):
+            assert read_json("rank-%s" % rank, tmp)["rank"] == rank