change: Add 2.0-gpu-py310-cu121 to unsupported dlc images for the estimator.

Add additional unsupported image tests.
Clean up tests.

Author: viclzhu

src/sagemaker/estimator.py

@@ -3198,6 +3198,7 @@ class Framework(EstimatorBase):
     """
 
     _framework_name = None
+    UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM = ("2.0.1-gpu-py310-cu121", "2.0-gpu-py310-cu121")
 
     def __init__(
         self,
@@ -3415,7 +3416,6 @@ def __init__(
         self.checkpoint_s3_uri = checkpoint_s3_uri
         self.checkpoint_local_path = checkpoint_local_path
         self.enable_sagemaker_metrics = enable_sagemaker_metrics
-        self.unsupported_dlc_image_for_sm_parallelism = ["2.0.1-gpu-py310-cu121"]
 
     def _prepare_for_training(self, job_name=None):
         """Set hyperparameters needed for training. This method will also validate ``source_dir``.
@@ -3854,7 +3854,7 @@ def _distribution_configuration(self, distribution):
             smdataparallel_enabled = smdistributed.get("dataparallel", {}).get("enabled", False)
             p5_enabled = bool("p5.48xlarge" in self.instance_type)
             img_uri = "" if self.image_uri is None else self.image_uri
-            for unsupported_image in self.unsupported_dlc_image_for_sm_parallelism:
+            for unsupported_image in Framework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
                 if unsupported_image in img_uri and not torch_distributed_enabled: #disabling DLC images with CUDA12
                     raise ValueError(f"SMDistributed is currently incompatible with DLC image: {img_uri}. (Could be due to CUDA version being greater than 11.)")
             if not torch_distributed_enabled and p5_enabled: #disabling p5 when torch distributed is disabled

tests/unit/test_estimator.py

@@ -322,35 +322,57 @@ def training_job_description(sagemaker_session):
     sagemaker_session.describe_training_job = mock_describe_training_job
     return returned_job_description
 
+
+def test_validate_smdistributed_unsupported_image_raises(sagemaker_session):
+    # Test unsupported image raises error.
+    for unsupported_image in DummyFramework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
+        # Fail due to unsupported CUDA12 DLC image.
+        f = DummyFramework(
+            "some_script.py",
+            role="DummyRole",
+            instance_type="ml.p4d.24xlarge",
+            sagemaker_session=sagemaker_session,
+            output_path="outputpath",
+            image_uri=unsupported_image,
+        )
+        with pytest.raises(ValueError):
+            f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED)
+        with pytest.raises(ValueError):
+            f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED)
+
+    # Test unsupported image with suffix raises error.
+    for unsupported_image in DummyFramework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
+        # Fail due to unsupported CUDA12 DLC image.
+        f = DummyFramework(
+            "some_script.py",
+            role="DummyRole",
+            instance_type="ml.p4d.24xlarge",
+            sagemaker_session=sagemaker_session,
+            output_path="outputpath",
+            image_uri=unsupported_image + "-ubuntu20.04-sagemaker-pr-3303",
+        )
+        with pytest.raises(ValueError):
+            f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED)
+        with pytest.raises(ValueError):
+            f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED)
+
+
 def test_validate_smdistributed_p5_raises(sagemaker_session):
-    # supported DLC image
+    # Supported DLC image.
     f = DummyFramework(
         "some_script.py", 
         role="DummyRole",
         instance_type="ml.p5.48xlarge", 
         sagemaker_session=sagemaker_session,
         output_path="outputpath",
-        image_uri="some_acceptable_image"
+        image_uri="some_acceptable_image",
     )
-    #both fail because instance type is p5 and torch_distributed is off
+    # Both fail because instance type is p5 and torch_distributed is off.
     with pytest.raises(ValueError):
         f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED) 
     with pytest.raises(ValueError):
         f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED) 
-    # unsupported DLC image
-    f = DummyFramework(
-        "some_script.py",
-        role="DummyRole",
-        instance_type="ml.p5.48xlarge",
-        sagemaker_session=sagemaker_session,
-        output_path="outputpath",
-        image_uri="ecr-url/2.0.1-gpu-py310-cu121-ubuntu20.04-sagemaker-pr-3303"
-    )
-    #both fail due to unsupported CUDA12 DLC image
-    with pytest.raises(ValueError):
-        f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED)
-    with pytest.raises(ValueError):
-        f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED)
+
 
 def test_validate_smdistributed_p5_not_raises(sagemaker_session):
     f = DummyFramework(
@@ -359,20 +381,23 @@ def test_validate_smdistributed_p5_not_raises(sagemaker_session):
         instance_type="ml.p5.48xlarge",
         sagemaker_session=sagemaker_session,
         output_path="outputpath",
-        image_uri="ecr-url/2.0.1-gpu-py310-cu121-ubuntu20.04-sagemaker-pr-3303"
+        image_uri="ecr-url/2.0.1-gpu-py310-cu121-ubuntu20.04-sagemaker-pr-3303",
     )
-    #testing with p5 instance and torch_distributed enabled
+    # Testing with p5 instance and torch_distributed enabled.
     f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED)
     f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED)
+
+
+def test_validate_smdistributed_backward_compat_p4_not_raises(sagemaker_session):
     f = DummyFramework(
         "some_script.py",
         role="DummyRole",
-        instance_type="ml.p4.24xlarge",
+        instance_type="ml.p4d.24xlarge",
         sagemaker_session=sagemaker_session,
         output_path="outputpath",
-        image_uri="some_acceptable_image"
+        image_uri="some_acceptable_image",
     )
-    #testing backwards compatability with p4d instances
+    # Testing backwards compatability with p4d instances.
     f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED)
     f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED)