fix: fix formatting

Author: viclzhu

src/sagemaker/estimator.py

@@ -3852,13 +3852,22 @@ def _distribution_configuration(self, distribution):
             torch_distributed_enabled = distribution.get("torch_distributed").get("enabled", False)
             smdistributed = distribution["smdistributed"]
             smdataparallel_enabled = smdistributed.get("dataparallel", {}).get("enabled", False)
-            p5_enabled = bool("p5.48xlarge" in self.instance_type)
+            p5_enabled = "p5.48xlarge" in self.instance_type
             img_uri = "" if self.image_uri is None else self.image_uri
             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
-                raise ValueError("SMModelParallel and SMDataParallel currently do not support p5 instances.")
+                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
+                raise ValueError(
+                    "SMModelParallel and SMDataParallel currently do not support p5 instances."
+                )
             # smdistributed strategy selected with supported instance type
             distribution_config[self.LAUNCH_SM_DDP_ENV_NAME] = smdataparallel_enabled
             distribution_config[self.INSTANCE_TYPE] = self.instance_type

tests/unit/test_estimator.py

@@ -170,19 +170,19 @@
 }
 DISTRIBUTION_SM_DDP_ENABLED = {
     "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "options"}},
-    "torch_distributed": {"enabled": False}
+    "torch_distributed": {"enabled": False},
 }
 DISTRIBUTION_SM_DDP_DISABLED = {
     "smdistributed": {"enabled": True},
-    "torch_distributed": {"enabled": False}
+    "torch_distributed": {"enabled": False},
 }
 DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED = {
     "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "options"}},
-    "torch_distributed": {"enabled": True}
+    "torch_distributed": {"enabled": True},
 }
 DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED = {
     "smdistributed": {"enabled": True},
-    "torch_distributed": {"enabled": True}
+    "torch_distributed": {"enabled": True},
 }
 MOCKED_S3_URI = "s3://mocked_s3_uri_from_source_dir"
 _DEFINITION_CONFIG = PipelineDefinitionConfig(use_custom_job_prefix=False)
@@ -360,18 +360,18 @@ def test_validate_smdistributed_unsupported_image_raises(sagemaker_session):
 def test_validate_smdistributed_p5_raises(sagemaker_session):
     # Supported DLC image.
     f = DummyFramework(
-        "some_script.py", 
+        "some_script.py",
         role="DummyRole",
-        instance_type="ml.p5.48xlarge", 
+        instance_type="ml.p5.48xlarge",
         sagemaker_session=sagemaker_session,
         output_path="outputpath",
         image_uri="some_acceptable_image",
     )
     # Both fail because instance type is p5 and torch_distributed is off.
     with pytest.raises(ValueError):
-        f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED) 
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED)
     with pytest.raises(ValueError):
-        f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED) 
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED)
 
 
 def test_validate_smdistributed_p5_not_raises(sagemaker_session):