Added unit and integration test

Author: karan6181

tests/integ/test_smdataparallel_pt.py

@@ -47,7 +47,9 @@ def test_smdataparallel_pt_mnist(
         sagemaker_session=sagemaker_session,
         framework_version=pytorch_training_latest_version,
         py_version=pytorch_training_latest_py_version,
-        distribution={"smdistributed": {"dataparallel": {"enabled": True}}},
+        distribution={
+            "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "--verbose"}}
+        },
     )
 
     with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):

tests/integ/test_smdataparallel_tf.py

@@ -47,7 +47,9 @@ def test_smdataparallel_tf_mnist(
         sagemaker_session=sagemaker_session,
         framework_version=tensorflow_training_latest_version,
         py_version=tensorflow_training_latest_py_version,
-        distribution={"smdistributed": {"dataparallel": {"enabled": True}}},
+        distribution={
+            "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "--verbose"}}
+        },
     )
 
     with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):

tests/unit/test_fw_utils.py

@@ -553,11 +553,15 @@ def test_validate_version_or_image_args_raises():
 
 def test_validate_smdistributed_not_raises():
     smdataparallel_enabled = {"smdistributed": {"dataparallel": {"enabled": True}}}
+    smdataparallel_enabled_custom_mpi = {
+        "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "--verbose"}}
+    }
     smdataparallel_disabled = {"smdistributed": {"dataparallel": {"enabled": False}}}
     instance_types = list(fw_utils.SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES)
 
     good_args = [
         (smdataparallel_enabled, "custom-container"),
+        (smdataparallel_enabled_custom_mpi, "custom-container"),
         (smdataparallel_disabled, "custom-container"),
     ]
     frameworks = ["tensorflow", "pytorch"]
@@ -576,17 +580,17 @@ def test_validate_smdistributed_not_raises():
 
 def test_validate_smdistributed_raises():
     bad_args = [
-        {"smdistributed": {"dataparallel": {"enabled": True}}},
         {"smdistributed": "dummy"},
         {"smdistributed": {"dummy"}},
         {"smdistributed": {"dummy": "val"}},
         {"smdistributed": {"dummy": {"enabled": True}}},
     ]
+    instance_types = list(fw_utils.SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES)
     frameworks = ["tensorflow", "pytorch"]
-    for framework, distribution in product(frameworks, bad_args):
+    for framework, distribution, instance_type in product(frameworks, bad_args, instance_types):
         with pytest.raises(ValueError):
             fw_utils.validate_smdistributed(
-                instance_type=None,
+                instance_type=instance_type,
                 framework_name=framework,
                 framework_version=None,
                 py_version=None,
@@ -624,6 +628,9 @@ def test_validate_smdataparallel_args_raises():
 
 def test_validate_smdataparallel_args_not_raises():
     smdataparallel_enabled = {"smdistributed": {"dataparallel": {"enabled": True}}}
+    smdataparallel_enabled_custom_mpi = {
+        "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "--verbose"}}
+    }
     smdataparallel_disabled = {"smdistributed": {"dataparallel": {"enabled": False}}}
 
     # Cases {PT|TF2}
@@ -643,6 +650,8 @@ def test_validate_smdataparallel_args_not_raises():
         ("ml.p3.16xlarge", "pytorch", "1.7", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.8.0", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.8", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py3", smdataparallel_enabled_custom_mpi),
+        ("ml.p3.16xlarge", "pytorch", "1.8.0", "py3", smdataparallel_enabled_custom_mpi),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(