chore(codestyle): fix black formatting and unused import

sagemaker-core/src/sagemaker/core/workflow/job_config_document.py

@@ -28,6 +28,7 @@
   number or boolean reaches the service quoted. Parameterise string-valued
   fields only.
 """
+
 from __future__ import absolute_import
 
 import json
@@ -67,9 +68,7 @@ def _placeholder(obj):
         if isinstance(obj, PipelineVariable):
             variables.append(obj)
             return "%s%d__" % (token_prefix, len(variables) - 1)
-        raise TypeError(
-            "Object of type %s is not JSON serializable" % obj.__class__.__name__
-        )
+        raise TypeError("Object of type %s is not JSON serializable" % obj.__class__.__name__)
 
     document = json.dumps(job_config, default=_placeholder)
     if not variables:

sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Shared capture paths for producers running under a ``PipelineSession``."""
+
 from __future__ import absolute_import
 
 from typing import Any, Dict, List, Optional
@@ -87,9 +88,11 @@ def capture_training_request(pipeline_session, create_args: Dict[str, Any]) -> N
         # Dict-form tags arrive in both shapes: lowercase keys from the SDK's own
         # tag builders, PascalCase from callers passing wire-form dicts through.
         request["Tags"] = [
-            Tag(key=tag.get("key", tag.get("Key")), value=tag.get("value", tag.get("Value")))
-            if isinstance(tag, dict)
-            else tag
+            (
+                Tag(key=tag.get("key", tag.get("Key")), value=tag.get("value", tag.get("Value")))
+                if isinstance(tag, dict)
+                else tag
+            )
             for tag in request["Tags"]
         ]
     pipeline_session._intercept_create_request(serialize(request), None, "train")

sagemaker-core/tests/unit/workflow/test_job_config_document.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Tests for the ``JobConfigDocument`` encoder."""
+
 from __future__ import absolute_import
 
 import json

sagemaker-core/tests/unit/workflow/test_pipeline_capture.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Tests for the shared ``PipelineSession`` capture helpers."""
+
 from __future__ import absolute_import
 
 import json

sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py

@@ -19,6 +19,7 @@
 -- `ProcessingStep` expects `"run"` from here while interception is called with
 `"process"` -- so a step guard cannot be satisfied by the interception name alone.
 """
+
 from __future__ import absolute_import
 
 import pytest

sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py

@@ -871,7 +871,9 @@ def arguments(self) -> RequestType:
         """
         from sagemaker.core.workflow.execution_variables import ExecutionVariables
         from sagemaker.core.workflow.functions import Join
-        from sagemaker.core.workflow.job_config_document import convert_job_config_document_to_string
+        from sagemaker.core.workflow.job_config_document import (
+            convert_job_config_document_to_string,
+        )
         from sagemaker.core.workflow.utilities import execute_job_functions
         from sagemaker.core.workflow.utilities import _pipeline_config
 

sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py

@@ -14,6 +14,7 @@
 resolves the execution-scoped output path inside `JobConfigDocument`, and
 creates the AgentRFT job from it. The execution is stopped once the job exists.
 """
+
 from __future__ import absolute_import
 
 import json

sagemaker-mlops/tests/unit/workflow/test_job_step.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Unit tests for `StepTypeEnum.JOB` and `JobStep`."""
+
 from __future__ import absolute_import
 
 import pytest

sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py

@@ -2,7 +2,7 @@
 
 import json
 import warnings
-from unittest.mock import MagicMock, patch, PropertyMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -1057,9 +1057,10 @@ def test_declaration_wins_over_the_model_trainer_branch(self):
     def test_train_returns_step_arguments_and_submits_nothing(self):
         session = self._pipeline_session()
         trainer = self._make_trainer(sagemaker_session=session)
-        with self._pinned_name(), patch(
-            "sagemaker.train.multi_turn_rl_trainer.Job.create"
-        ) as mock_create:
+        with (
+            self._pinned_name(),
+            patch("sagemaker.train.multi_turn_rl_trainer.Job.create") as mock_create,
+        ):
             step_args = trainer.train()
         assert isinstance(step_args, _StepArguments)
         assert step_args.caller_name == "create_job"
@@ -1137,9 +1138,7 @@ def test_wait_is_overridden_and_announced(self):
             warnings.simplefilter("always")
             step_args = trainer.train(wait=True)
         assert step_args.func_kwargs.get("wait") is False
-        assert any(
-            "No Wait" in str(w.message) for w in caught
-        ), [str(w.message) for w in caught]
+        assert any("No Wait" in str(w.message) for w in caught), [str(w.message) for w in caught]
 
     # --- parity with the direct submission path -------------------------------
 
@@ -1151,11 +1150,13 @@ def _record(**kwargs):
             recorded.update(kwargs)
             return MagicMock()
 
-        with self._pinned_name(), patch(
-            "sagemaker.train.multi_turn_rl_trainer.Job.create", side_effect=_record
-        ), patch(
-            "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job",
-            return_value=MagicMock(),
+        with (
+            self._pinned_name(),
+            patch("sagemaker.train.multi_turn_rl_trainer.Job.create", side_effect=_record),
+            patch(
+                "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job",
+                return_value=MagicMock(),
+            ),
         ):
             trainer.train(wait=False)
         return recorded
@@ -1203,9 +1204,9 @@ def test_captured_request_matches_what_job_create_would_send(self):
         }
 
         session = self._pipeline_session()
-        captured = dict(self._capture(
-            self._make_trainer(sagemaker_session=session, tags=tags), session
-        ).args)
+        captured = dict(
+            self._capture(self._make_trainer(sagemaker_session=session, tags=tags), session).args
+        )
         # The document is captured as the raw config dict for JobStep to scope and
         # encode; compare it to the direct route's config, and the rest to the wire.
         document = captured.pop("JobConfigDocument")
@@ -1221,9 +1222,9 @@ def test_captured_request_matches_job_create_with_both_tag_forms(self):
             self._make_trainer(sagemaker_session=self._direct_session(), tags=tags)
         )
         session = self._pipeline_session()
-        captured = dict(self._capture(
-            self._make_trainer(sagemaker_session=session, tags=tags), session
-        ).args)
+        captured = dict(
+            self._capture(self._make_trainer(sagemaker_session=session, tags=tags), session).args
+        )
         document = captured.pop("JobConfigDocument")
         assert document == json.loads(direct["job_config_document"])
         wire = self._job_create_wire_request(direct, direct["job_config_document"])
@@ -1256,7 +1257,9 @@ def test_document_describes_the_same_config_on_both_routes(self):
         session = self._pipeline_session()
         trainer = self._make_trainer(sagemaker_session=session)
         captured = self._capture(trainer, session).args["JobConfigDocument"]
-        direct = self._make_trainer(sagemaker_session=self._direct_session())._build_job_config_document()
+        direct = self._make_trainer(
+            sagemaker_session=self._direct_session()
+        )._build_job_config_document()
         assert isinstance(captured, dict)
         assert captured == json.loads(direct)
 
@@ -1281,7 +1284,8 @@ def test_pipeline_variable_in_the_job_config_survives_capture(self):
     def test_pipeline_variable_document_would_be_a_type_error_unencoded(self):
         """Locks the reason the encoder is needed rather than assuming it."""
         trainer = self._make_trainer(
-            sagemaker_session=self._direct_session(), s3_output_path=ParameterString(name="OutputPath")
+            sagemaker_session=self._direct_session(),
+            s3_output_path=ParameterString(name="OutputPath"),
         )
         with pytest.raises(TypeError, match="not JSON serializable"):
             trainer._build_job_config_document()
@@ -1292,11 +1296,13 @@ def test_direct_path_still_submits_and_returns_a_job(self):
         """Criterion: a normal session behaves exactly as before."""
         job_handle = MagicMock()
         trainer = self._make_trainer(sagemaker_session=self._direct_session())
-        with self._pinned_name(), patch(
-            "sagemaker.train.multi_turn_rl_trainer.Job.create"
-        ) as mock_create, patch(
-            "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job",
-            return_value=job_handle,
+        with (
+            self._pinned_name(),
+            patch("sagemaker.train.multi_turn_rl_trainer.Job.create") as mock_create,
+            patch(
+                "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job",
+                return_value=job_handle,
+            ),
         ):
             returned = trainer.train(wait=False)
         mock_create.assert_called_once()
@@ -1307,11 +1313,13 @@ def test_direct_path_still_submits_and_returns_a_job(self):
     def test_direct_path_honours_wait(self):
         job_handle = MagicMock()
         trainer = self._make_trainer(sagemaker_session=self._direct_session())
-        with self._pinned_name(), patch(
-            "sagemaker.train.multi_turn_rl_trainer.Job.create"
-        ), patch(
-            "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job",
-            return_value=job_handle,
+        with (
+            self._pinned_name(),
+            patch("sagemaker.train.multi_turn_rl_trainer.Job.create"),
+            patch(
+                "sagemaker.train.multi_turn_rl_trainer.AgentRFTJob.from_job",
+                return_value=job_handle,
+            ),
         ):
             trainer.train(wait=True)
         job_handle.wait.assert_called_once()