From 2a3cf9840d8f861e8b2a5bd235c94b6579c7e4c8 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Wed, 7 Oct 2026 16:25:37 -0700 Subject: [PATCH] chore(codestyle): fix black formatting and unused import Apply black (26.3.1, line-length 100) to files that were committed unformatted, and remove an unused import, fixing the failing codestyle-doc-tests (black-check + flake8) on sagemaker-core, sagemaker-mlops, and sagemaker-train. - sagemaker-core: reformat workflow/job_config_document.py, workflow/pipeline_capture.py and three workflow unit tests. - sagemaker-mlops: reformat workflow/steps.py and two workflow tests. - sagemaker-train: reformat test_multi_turn_rl_trainer.py and drop the unused unittest.mock.PropertyMock import (flake8 F401). Formatting-only (black preserves the AST); the sole semantic change is the removal of the unused import. --- .../core/workflow/job_config_document.py | 5 +- .../core/workflow/pipeline_capture.py | 9 ++- .../unit/workflow/test_job_config_document.py | 1 + .../unit/workflow/test_pipeline_capture.py | 1 + .../workflow/test_retrieve_caller_name.py | 1 + .../src/sagemaker/mlops/workflow/steps.py | 4 +- .../integ/workflow/test_job_step_mtrl.py | 1 + .../tests/unit/workflow/test_job_step.py | 1 + .../unit/train/test_multi_turn_rl_trainer.py | 68 +++++++++++-------- 9 files changed, 54 insertions(+), 37 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py b/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py index 63ed77eee5..09fe329f98 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py +++ b/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: diff --git a/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py b/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py index e48b059c68..8e319b64fb 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py +++ b/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") diff --git a/sagemaker-core/tests/unit/workflow/test_job_config_document.py b/sagemaker-core/tests/unit/workflow/test_job_config_document.py index 5ad36a8746..ef3ffec949 100644 --- a/sagemaker-core/tests/unit/workflow/test_job_config_document.py +++ b/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 diff --git a/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py b/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py index 3ad2849cae..5e7867682c 100644 --- a/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py +++ b/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 diff --git a/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py b/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py index ff08fa1d50..5de23c6fd9 100644 --- a/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py +++ b/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 diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 9c25a35310..333f32616c 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/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 diff --git a/sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py b/sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py index 1fd4ac646b..8931365205 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_job_step_mtrl.py +++ b/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 diff --git a/sagemaker-mlops/tests/unit/workflow/test_job_step.py b/sagemaker-mlops/tests/unit/workflow/test_job_step.py index ac9dcb28b1..4e6e572409 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_job_step.py +++ b/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 diff --git a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py b/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py index e23f615d81..40592437f3 100644 --- a/sagemaker-train/tests/unit/train/test_multi_turn_rl_trainer.py +++ b/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()