fix unit tests

Author: tombousso

tests/unit/test_fw_utils.py

@@ -853,18 +853,12 @@ def test_validate_smdataparallel_args_raises():
     smdataparallel_enabled = {"smdistributed": {"dataparallel": {"enabled": True}}}
 
     # Cases {PT|TF2}
-    # 1. None instance type
-    # 2. incorrect instance type
-    # 3. incorrect python version
-    # 4. incorrect framework version
+    # 1. incorrect python version
+    # 2. incorrect framework version
 
     bad_args = [
-        (None, "tensorflow", "2.3.1", "py3", smdataparallel_enabled),
-        ("ml.p3.2xlarge", "tensorflow", "2.3.1", "py3", smdataparallel_enabled),
         ("ml.p3dn.24xlarge", "tensorflow", "2.3.1", "py2", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "1.3.1", "py3", smdataparallel_enabled),
-        (None, "pytorch", "1.6.0", "py3", smdataparallel_enabled),
-        ("ml.p3.2xlarge", "pytorch", "1.6.0", "py3", smdataparallel_enabled),
         ("ml.p3dn.24xlarge", "pytorch", "1.6.0", "py2", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.5.0", "py3", smdataparallel_enabled),
     ]
@@ -966,74 +960,6 @@ def test_validate_smdataparallel_args_not_raises():
         )
 
 
-def test_validate_pytorchddp_not_raises():
-    # Case 1: Framework is not PyTorch
-    fw_utils.validate_pytorch_distribution(
-        distribution=None,
-        framework_name="tensorflow",
-        framework_version="2.9.1",
-        py_version="py3",
-        image_uri="custom-container",
-    )
-    # Case 2: Framework is PyTorch, but distribution is not PyTorchDDP
-    pytorchddp_disabled = {"pytorchddp": {"enabled": False}}
-    fw_utils.validate_pytorch_distribution(
-        distribution=pytorchddp_disabled,
-        framework_name="pytorch",
-        framework_version="1.10",
-        py_version="py3",
-        image_uri="custom-container",
-    )
-    # Case 3: Framework is PyTorch, Distribution is PyTorchDDP enabled, supported framework and py versions
-    pytorchddp_enabled = {"pytorchddp": {"enabled": True}}
-    pytorchddp_supported_fw_versions = [
-        "1.10",
-        "1.10.0",
-        "1.10.2",
-        "1.11",
-        "1.11.0",
-        "1.12",
-        "1.12.0",
-        "1.12.1",
-        "1.13.1",
-        "2.0.0",
-        "2.0.1",
-        "2.1.0",
-        "2.2.0",
-    ]
-    for framework_version in pytorchddp_supported_fw_versions:
-        fw_utils.validate_pytorch_distribution(
-            distribution=pytorchddp_enabled,
-            framework_name="pytorch",
-            framework_version=framework_version,
-            py_version="py3",
-            image_uri="custom-container",
-        )
-
-
-def test_validate_pytorchddp_raises():
-    pytorchddp_enabled = {"pytorchddp": {"enabled": True}}
-    # Case 1: Unsupported framework version
-    with pytest.raises(ValueError):
-        fw_utils.validate_pytorch_distribution(
-            distribution=pytorchddp_enabled,
-            framework_name="pytorch",
-            framework_version="1.8",
-            py_version="py3",
-            image_uri=None,
-        )
-
-    # Case 2: Unsupported Py version
-    with pytest.raises(ValueError):
-        fw_utils.validate_pytorch_distribution(
-            distribution=pytorchddp_enabled,
-            framework_name="pytorch",
-            framework_version="1.10",
-            py_version="py2",
-            image_uri=None,
-        )
-
-
 def test_validate_torch_distributed_not_raises():
     # Case 1: Framework is PyTorch, but torch_distributed is not enabled
     torch_distributed_disabled = {"torch_distributed": {"enabled": False}}

tests/unit/test_pytorch.py

@@ -801,14 +801,15 @@ def test_pytorch_ddp_distribution_configuration(
         distribution=pytorch.distribution
     )
     expected_torch_ddp = {
-        "sagemaker_pytorch_ddp_enabled": True,
+        "sagemaker_distributed_dataparallel_enabled": True,
+        "sagemaker_distributed_dataparallel_custom_mpi_options": "",
         "sagemaker_instance_type": test_instance_type,
     }
     assert actual_pytorch_ddp == expected_torch_ddp
 
 
 def test_pytorch_ddp_distribution_configuration_unsupported(sagemaker_session):
-    unsupported_framework_version = "1.9.1"
+    unsupported_framework_version = "1.5.0"
     unsupported_py_version = "py2"
     with pytest.raises(ValueError) as error:
         _pytorch_estimator(