fix formatting

Author: jerrypeng7773

src/sagemaker/estimator.py

@@ -77,7 +77,7 @@
 from sagemaker.workflow.pipeline_context import (
     PipelineSession,
     runnable_by_pipeline,
-    is_pipeline_entities
+    is_pipeline_entities,
 )
 
 logger = logging.getLogger(__name__)
@@ -1335,7 +1335,10 @@ def register(
     @property
     def model_data(self):
         """str: The model location in S3. Only set if Estimator has been ``fit()``."""
-        if self.latest_training_job is not None and type(self.sagemaker_session) is not PipelineSession:
+        if (
+            self.latest_training_job is not None
+            and type(self.sagemaker_session) is not PipelineSession
+        ):
             model_uri = self.sagemaker_session.sagemaker_client.describe_training_job(
                 TrainingJobName=self.latest_training_job.name
             )["ModelArtifacts"]["S3ModelArtifacts"]
@@ -1762,7 +1765,7 @@ def start_new(cls, estimator, inputs, experiment_config):
         """
         train_args = cls._get_train_args(estimator, inputs, experiment_config)
         if type(estimator.sagemaker_session) is PipelineSession:
-            train_args['pipeline_session'] = estimator.sagemaker_session
+            train_args["pipeline_session"] = estimator.sagemaker_session
 
         estimator.sagemaker_session.train(**train_args)
 

src/sagemaker/processing.py

@@ -1759,4 +1759,3 @@ def _set_entrypoint(self, command, user_script_name):
             )
         )
         self.entrypoint = self.framework_entrypoint_command + [user_script_location]
-

src/sagemaker/transformer.py

@@ -17,7 +17,11 @@
 
 from sagemaker.job import _Job
 from sagemaker.session import Session
-from sagemaker.workflow.pipeline_context import PipelineSession, runnable_by_pipeline, is_pipeline_entities
+from sagemaker.workflow.pipeline_context import (
+    PipelineSession,
+    runnable_by_pipeline,
+    is_pipeline_entities,
+)
 from sagemaker.utils import base_name_from_image, name_from_base
 
 
@@ -375,7 +379,7 @@ def start_new(
             model_client_config,
         )
         if type(transformer.sagemaker_session) is PipelineSession:
-            transform_args['pipeline_session'] = transformer.sagemaker_session
+            transform_args["pipeline_session"] = transformer.sagemaker_session
         transformer.sagemaker_session.transform(**transform_args)
 
         return cls(transformer.sagemaker_session, transformer._current_job_name)

src/sagemaker/tuner.py

@@ -41,7 +41,7 @@
 from sagemaker.workflow.pipeline_context import (
     PipelineSession,
     runnable_by_pipeline,
-    is_pipeline_entities
+    is_pipeline_entities,
 )
 
 from sagemaker.session import Session
@@ -366,9 +366,7 @@ def _prepare_static_hyperparameters(
         """Prepare static hyperparameters for one estimator before tuning."""
         # Remove any hyperparameter that will be tuned
         static_hyperparameters = {
-            str(k): str(v)
-            if not is_pipeline_entities(v)
-            else v
+            str(k): str(v) if not is_pipeline_entities(v) else v
             for (k, v) in estimator.hyperparameters().items()
         }
         for hyperparameter_name in hyperparameter_ranges.keys():
@@ -475,7 +473,7 @@ def _fit_with_estimator_dict(self, inputs, job_name, include_cls_metadata, estim
         self._validate_dict_argument(
             name="include_cls_metadata",
             value=include_cls_metadata if include_cls_metadata else {},
-            allowed_keys=estimator_names
+            allowed_keys=estimator_names,
         )
         self._validate_dict_argument(
             name="estimator_kwargs", value=estimator_kwargs, allowed_keys=estimator_names
@@ -1480,7 +1478,7 @@ def start_new(cls, tuner, inputs):
         """
         tuner_args = cls._get_tuner_args(tuner, inputs)
         if type(tuner.sagemaker_session) is PipelineSession:
-            tuner_args['pipeline_session'] = PipelineSession
+            tuner_args["pipeline_session"] = PipelineSession
 
         tuner.sagemaker_session.create_tuning_job(**tuner_args)
 

src/sagemaker/workflow/pipeline.py

@@ -41,6 +41,7 @@
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.utilities import list_to_request
 
+
 @attr.s
 class Pipeline(Entity):
     """Pipeline for workflow.

src/sagemaker/workflow/pipeline_context.py

@@ -1,3 +1,16 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The pipeline context for workflow"""
 from __future__ import absolute_import
 
 import warnings
@@ -11,20 +24,20 @@
 
 
 class PipelineSession(Session):
-    """Managing interactions with the Amazon SageMaker APIs and any other AWS services needed
-            under SageMaker Model-Building Pipeline Context
+    """Managing interactions with SageMaker APIs and AWS services needed under SageMaker Model-Building Pipeline Context
+
+    This class inherits the SageMaker session, it provides convenient methods for manipulating entities
+    and resources that Amazon SageMaker uses, such as training jobs, endpoints, and input datasets in S3.
+    When composing SageMaker Model-Building Pipeline, PipelineSession is recommended over
+    regular SageMaker Session
+    """
 
-        This class inherits the SageMaker session, it provides convenient methods for manipulating entities
-        and resources that Amazon SageMaker uses, such as training jobs, endpoints, and input datasets in S3.
-        When composing SageMaker Model-Building Pipeline, PipelineSession is recommended over
-        regular SageMaker Session
-        """
     def __init__(
-            self,
-            boto_session=None,
-            sagemaker_client=None,
-            default_bucket=None,
-            settings=SessionSettings(),
+        self,
+        boto_session=None,
+        sagemaker_client=None,
+        default_bucket=None,
+        settings=SessionSettings(),
     ):
         """Initialize a ``PipelineSession``.
 
@@ -66,16 +79,17 @@ def context(self, args: Dict):
 def runnable_by_pipeline(run_func):
     """A convenient Decorator
 
-        This is a decorator designed to annotate, during pipeline session, the methods that downstream managed to
-            1. preprocess user inputs, outputs, and configurations
-            2. generate the create request
-            3. start the job.
-        For instance, `Processor.run`, `Estimator.fit`, or `Transformer.transform`. This decorator will
-        essentially run 1, and capture the request shape from 2, then instead of starting a new job in 3, it will
-        return request shape from 2 to `sagemaker.workflow.steps.Step`. The request shape will be used to construct
-        the arguments needed to compose that particular step as part of the pipeline. The job will be started during
-        pipeline execution.
+    This is a decorator designed to annotate, during pipeline session, the methods that downstream managed to
+        1. preprocess user inputs, outputs, and configurations
+        2. generate the create request
+        3. start the job.
+    For instance, `Processor.run`, `Estimator.fit`, or `Transformer.transform`. This decorator will
+    essentially run 1, and capture the request shape from 2, then instead of starting a new job in 3, it will
+    return request shape from 2 to `sagemaker.workflow.steps.Step`. The request shape will be used to construct
+    the arguments needed to compose that particular step as part of the pipeline. The job will be started during
+    pipeline execution.
     """
+
     def wrapper(*args, **kwargs):
         if type(args[0].sagemaker_session) is PipelineSession:
             run_func_sig = inspect.signature(run_func)
@@ -85,29 +99,30 @@ def wrapper(*args, **kwargs):
             for i, (arg_name, param) in enumerate(run_func_sig.parameters.items()):
                 if i >= len(arg_list):
                     break
-                if arg_name == 'wait':
+                if arg_name == "wait":
                     override_wait = True
                     arg_list[i] = False
-                elif arg_name == 'logs':
+                elif arg_name == "logs":
                     override_logs = True
                     arg_list[i] = False
 
             args = tuple(arg_list)
 
             if not override_wait:
-                kwargs['wait'] = False
+                kwargs["wait"] = False
             if not override_logs:
-                kwargs['logs'] = False
+                kwargs["logs"] = False
 
             warnings.warn(
                 "Running within a PipelineSession, there will be No Wait, "
                 "No Logs, and No Job being started.",
-                UserWarning
+                UserWarning,
             )
             run_func(*args, **kwargs)
             return args[0].sagemaker_session.context
         else:
             run_func(*args, **kwargs)
+
     return wrapper
 
 
@@ -119,4 +134,4 @@ def is_pipeline_entities(obj: Any) -> bool:
     Returns:
         bool: if the given object is a pipeline Parameter, Expression, or Properties
     """
-    return isinstance(obj, (Parameter, Expression, Properties))
+    return isinstance(obj, (Parameter, Expression, Properties))

src/sagemaker/workflow/steps.py

@@ -285,8 +285,9 @@ def __init__(
         self.cache_config = cache_config
 
         if self.cache_config:
-            if (self.run_args and "ProfilerConfig" in self.run_args) or \
-                    (self.estimator is not None and not self.estimator.disable_profiler):
+            if (self.run_args and "ProfilerConfig" in self.run_args) or (
+                self.estimator is not None and not self.estimator.disable_profiler
+            ):
                 msg = (
                     "Profiling is enabled on the provided estimator. "
                     "The default profiler rule includes a timestamp "
@@ -299,10 +300,10 @@ def __init__(
         if not self.run_args:
             warnings.warn(
                 (
-                    "We are deprecating the instantiation of TrainingStep using \"estimator\"."
-                    "Instead, simply using \"run_args\"."
+                    'We are deprecating the instantiation of TrainingStep using "estimator".'
+                    'Instead, simply using "run_args".'
                 ),
-                DeprecationWarning
+                DeprecationWarning,
             )
 
     @property
@@ -464,10 +465,10 @@ def __init__(
                 raise ValueError("Inputs can't be None when transformer is given.")
             warnings.warn(
                 (
-                    "We are deprecating the instantiation of TransformStep using \"transformer\"."
-                    "Instead, simply using \"run_args\"."
+                    'We are deprecating the instantiation of TransformStep using "transformer".'
+                    'Instead, simply using "run_args".'
                 ),
-                DeprecationWarning
+                DeprecationWarning,
             )
 
     @property
@@ -493,7 +494,9 @@ def arguments(self) -> RequestType:
                 model_client_config=self.inputs.model_client_config,
                 experiment_config=dict(),
             )
-            request_dict = self.transformer.sagemaker_session._get_transform_request(**transform_args)
+            request_dict = self.transformer.sagemaker_session._get_transform_request(
+                **transform_args
+            )
 
         request_dict.pop("TransformJobName", None)
         return request_dict
@@ -598,10 +601,10 @@ def __init__(
 
             warnings.warn(
                 (
-                    "We are deprecating the instantiation of ProcessingStep using \"processor\"."
-                    "Instead, simply using \"run_args\"."
+                    'We are deprecating the instantiation of ProcessingStep using "processor".'
+                    'Instead, simply using "run_args".'
                 ),
-                DeprecationWarning
+                DeprecationWarning,
             )
 
     @property
@@ -738,10 +741,10 @@ def __init__(
         if not self.run_args:
             warnings.warn(
                 (
-                    "We are deprecating the instantiation of TuningStep using \"tuner\"."
-                    "Instead, simply using \"run_args\"."
+                    'We are deprecating the instantiation of TuningStep using "tuner".'
+                    'Instead, simply using "run_args".'
                 ),
-                DeprecationWarning
+                DeprecationWarning,
             )
 
     @property
@@ -808,5 +811,3 @@ def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = "") ->
                 "output/model.tar.gz",
             ],
         )
-
-

src/sagemaker/workflow/utilities.py

@@ -58,4 +58,3 @@ def hash_file(path: str) -> str:
             md5.update(data)
 
     return md5.hexdigest()
-

tests/unit/sagemaker/workflow/test_pipeline_session.py

@@ -33,5 +33,3 @@ def test_pipeline_session_init(sagemaker_client_config, boto_session):
     assert sess.sagemaker_client is not None
     assert sess.default_bucket() is not None
     assert sess.context is None
-
-