linting

Author: Lokiiiiii

src/sagemaker/estimator.py

@@ -2751,16 +2751,26 @@ def _validate_and_set_debugger_configs(self):
     def _validate_mwms_config(self):
         """Validate Multi Worker Mirrored Strategy configuration."""
         minimum_supported_framework_version = {
-                                                'tensorflow': {'framework_version': '2.9'},
-                                            }
+            "tensorflow": {"framework_version": "2.9"},
+        }
         if self._framework_name in minimum_supported_framework_version:
             for version_argument in minimum_supported_framework_version[self._framework_name]:
                 current = getattr(self, version_argument)
-                threshold = minimum_supported_framework_version[self._framework_name][version_argument]
+                threshold = minimum_supported_framework_version[self._framework_name][
+                    version_argument
+                ]
                 if Version(current) in SpecifierSet(f"< {threshold}"):
-                    raise ValueError("Multi Worker Mirrored Strategy is only supported from {} {} but received {}".format(version_argument, threshold, current))
+                    raise ValueError(
+                        "Multi Worker Mirrored Strategy is only supported from {} {} but received {}".format(
+                            version_argument, threshold, current
+                        )
+                    )
         else:
-            raise ValueError("Multi Worker Mirrored Strategy is currently only supported with {} frameworks but received {}".format(minimum_supported_framework_version.keys(), self._framework_name))
+            raise ValueError(
+                "Multi Worker Mirrored Strategy is currently only supported with {} frameworks but received {}".format(
+                    minimum_supported_framework_version.keys(), self._framework_name
+                )
+            )
 
     def _model_source_dir(self):
         """Get the appropriate value to pass as ``source_dir`` to a model constructor.

src/sagemaker/tensorflow/training_compiler/config.py

@@ -111,7 +111,9 @@ def validate(
                 raise ValueError(error_helper_string)
 
         if estimator.distribution and "multi_worker_mirrored_strategy" in estimator.distribution:
-            mwms_enabled = estimator.distribution.get("multi_worker_mirrored_strategy").get("enabled", False)
+            mwms_enabled = estimator.distribution.get("multi_worker_mirrored_strategy").get(
+                "enabled", False
+            )
             if mwms_enabled:
                 raise ValueError(
                     "Multi Worker Mirrored Strategy distributed training configuration "

tests/unit/sagemaker/tensorflow/test_estimator.py

@@ -571,8 +571,8 @@ def test_fit_mwms_unsupported(time, strftime, sagemaker_session):
         inputs = "s3://mybucket/train"
         tf.fit(inputs=inputs)
 
-    assert 'only supported from' in str(error)
-    assert 'but received' in str(error)
+    assert "only supported from" in str(error)
+    assert "but received" in str(error)
 
 
 def test_hyperparameters_no_model_dir(

tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

@@ -226,7 +226,7 @@ def test_mwms(
                 framework_version=tensorflow_training_version,
                 enable_sagemaker_metrics=False,
                 compiler_config=TrainingCompilerConfig(),
-                distribution={'multi_worker_mirrored_strategy': True},
+                distribution={"multi_worker_mirrored_strategy": True},
             ).fit()
 
     def test_python_2(

tests/unit/test_estimator.py

@@ -3521,9 +3521,9 @@ def test_mwms_distribution_configuration(sagemaker_session):
     )
     with pytest.raises(ValueError) as error:
         framework._distribution_configuration(distribution=DISTRIBUTION_MWMS_ENABLED)
-    
-    assert 'only supported with' in str(error)
-    assert 'but received' in str(error)
+
+    assert "only supported with" in str(error)
+    assert "but received" in str(error)
 
 
 def test_image_name_map(sagemaker_session):