Implement subclass compatibility for workflow pipeline job steps

Author: jerrypeng7773

src/sagemaker/amazon/amazon_estimator.py

@@ -27,6 +27,7 @@
 from sagemaker.estimator import EstimatorBase, _TrainingJob
 from sagemaker.inputs import FileSystemInput, TrainingInput
 from sagemaker.utils import sagemaker_timestamp
+from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 
 logger = logging.getLogger(__name__)
 
@@ -192,6 +193,7 @@ def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):
         self.feature_dim = feature_dim
         self.mini_batch_size = mini_batch_size
 
+    @runnable_by_pipeline
     def fit(
         self,
         records,

src/sagemaker/clarify.py

@@ -803,7 +803,8 @@ def _run(
                 output_name="analysis_result",
                 s3_upload_mode="EndOfJob",
             )
-            super().run(
+
+            return super().run(
                 inputs=[data_input, config_input],
                 outputs=[result_output],
                 wait=wait,
@@ -871,7 +872,7 @@ def run_pre_training_bias(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Pretraining-Bias")
-        self._run(
+        return self._run(
             data_config,
             analysis_config,
             wait,
@@ -957,7 +958,7 @@ def run_post_training_bias(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Posttraining-Bias")
-        self._run(
+        return self._run(
             data_config,
             analysis_config,
             wait,
@@ -1060,7 +1061,7 @@ def run_bias(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Bias")
-        self._run(
+        return self._run(
             data_config,
             analysis_config,
             wait,
@@ -1167,7 +1168,7 @@ def run_explainability(
                 job_name = utils.name_from_base(self.job_name_prefix)
             else:
                 job_name = utils.name_from_base("Clarify-Explainability")
-        self._run(
+        return self._run(
             data_config,
             analysis_config,
             wait,

src/sagemaker/estimator.py

@@ -73,9 +73,12 @@
     get_config_value,
     name_from_base,
 )
-from sagemaker.workflow.entities import Expression
-from sagemaker.workflow.parameters import Parameter
-from sagemaker.workflow.properties import Properties
+
+from sagemaker.workflow.pipeline_context import (
+    PipelineSession,
+    runnable_by_pipeline,
+    is_pipeline_entities
+)
 
 logger = logging.getLogger(__name__)
 
@@ -598,7 +601,7 @@ def _json_encode_hyperparameters(hyperparameters: Dict[str, Any]) -> Dict[str, A
         current_hyperparameters = hyperparameters
         if current_hyperparameters is not None:
             hyperparameters = {
-                str(k): (v if isinstance(v, (Parameter, Expression, Properties)) else json.dumps(v))
+                str(k): (v if is_pipeline_entities(v) else json.dumps(v))
                 for (k, v) in current_hyperparameters.items()
             }
         return hyperparameters
@@ -894,6 +897,7 @@ def latest_job_profiler_artifacts_path(self):
             )
         return None
 
+    @runnable_by_pipeline
     def fit(self, inputs=None, wait=True, logs="All", job_name=None, experiment_config=None):
         """Train a model using the input training dataset.
 
@@ -1331,7 +1335,7 @@ 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:
+        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"]
@@ -1757,6 +1761,9 @@ def start_new(cls, estimator, inputs, experiment_config):
             all information about the started training job.
         """
         train_args = cls._get_train_args(estimator, inputs, experiment_config)
+        if type(estimator.sagemaker_session) is PipelineSession:
+            train_args['pipeline_session'] = estimator.sagemaker_session
+
         estimator.sagemaker_session.train(**train_args)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
@@ -1801,7 +1808,7 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         current_hyperparameters = estimator.hyperparameters()
         if current_hyperparameters is not None:
             hyperparameters = {
-                str(k): (v if isinstance(v, (Parameter, Expression, Properties)) else str(v))
+                str(k): (v if is_pipeline_entities(v) else str(v))
                 for (k, v) in current_hyperparameters.items()
             }
 

src/sagemaker/processing.py

@@ -28,20 +28,21 @@
 
 from six.moves.urllib.parse import urlparse
 from six.moves.urllib.request import url2pathname
-
+import sagemaker
 from sagemaker import s3
 from sagemaker.job import _Job
 from sagemaker.local import LocalSession
 from sagemaker.utils import base_name_from_image, get_config_value, name_from_base
 from sagemaker.session import Session
-from sagemaker.workflow.properties import Properties
-from sagemaker.workflow.parameters import Parameter
-from sagemaker.workflow.entities import Expression
+from sagemaker.workflow.pipeline_context import (
+    PipelineSession,
+    runnable_by_pipeline,
+    is_pipeline_entities,
+)
 from sagemaker.dataset_definition.inputs import S3Input, DatasetDefinition
 from sagemaker.apiutils._base_types import ApiObject
 from sagemaker.s3 import S3Uploader
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -133,6 +134,7 @@ def __init__(
 
         self.sagemaker_session = sagemaker_session or Session()
 
+    @runnable_by_pipeline
     def run(
         self,
         inputs=None,
@@ -308,10 +310,10 @@ def _normalize_inputs(self, inputs=None, kms_key=None):
                 if file_input.input_name is None:
                     file_input.input_name = "input-{}".format(count)
 
-                if isinstance(file_input.source, Properties) or file_input.dataset_definition:
+                if is_pipeline_entities(file_input.source) or file_input.dataset_definition:
                     normalized_inputs.append(file_input)
                     continue
-                if isinstance(file_input.s3_input.s3_uri, (Parameter, Expression, Properties)):
+                if is_pipeline_entities(file_input.s3_input.s3_uri):
                     normalized_inputs.append(file_input)
                     continue
                 # If the source is a local path, upload it to S3
@@ -361,7 +363,7 @@ def _normalize_outputs(self, outputs=None):
                 # Generate a name for the ProcessingOutput if it doesn't have one.
                 if output.output_name is None:
                     output.output_name = "output-{}".format(count)
-                if isinstance(output.destination, (Parameter, Expression, Properties)):
+                if is_pipeline_entities(output.destination):
                     normalized_outputs.append(output)
                     continue
                 # If the output's destination is not an s3_uri, create one.
@@ -491,6 +493,7 @@ def get_run_args(
         """
         return RunArgs(code=code, inputs=inputs, outputs=outputs, arguments=arguments)
 
+    @runnable_by_pipeline
     def run(
         self,
         code,
@@ -765,6 +768,9 @@ def start_new(cls, processor, inputs, outputs, experiment_config):
         print("Inputs: ", process_args["inputs"])
         print("Outputs: ", process_args["output_config"]["Outputs"])
 
+        if type(processor.sagemaker_session) is PipelineSession:
+            process_args["pipeline_session"] = processor.sagemaker_session
+
         # Call sagemaker_session.process using the arguments dictionary.
         processor.sagemaker_session.process(**process_args)
 
@@ -1594,7 +1600,7 @@ def run(  # type: ignore[override]
         )
 
         # Submit a processing job.
-        super().run(
+        return super().run(
             code=s3_runproc_sh,
             inputs=inputs,
             outputs=outputs,
@@ -1753,3 +1759,4 @@ def _set_entrypoint(self, command, user_script_name):
             )
         )
         self.entrypoint = self.framework_entrypoint_command + [user_script_location]
+

src/sagemaker/session.py

@@ -466,6 +466,7 @@ def train(  # noqa: C901
         profiler_config=None,
         environment=None,
         retry_strategy=None,
+        pipeline_session=None,
     ):
         """Create an Amazon SageMaker training job.
 
@@ -551,7 +552,7 @@ def train(  # noqa: C901
             retry_strategy(dict): Defines RetryStrategy for InternalServerFailures.
                 * max_retry_attsmpts (int): Number of times a job should be retried.
                 The key in RetryStrategy is 'MaxRetryAttempts'.
-
+            pipeline_session (`sagemaker.workflow.pipeline_session.PipelineSession`) A pipeline session
         Returns:
             str: ARN of the training job, if it is created.
         """
@@ -585,6 +586,10 @@ def train(  # noqa: C901
             environment=environment,
             retry_strategy=retry_strategy,
         )
+        if pipeline_session:
+            pipeline_session.context = train_request
+            return
+
         LOGGER.info("Creating training-job with name: %s", job_name)
         LOGGER.debug("train request: %s", json.dumps(train_request, indent=4))
         self.sagemaker_client.create_training_job(**train_request)
@@ -856,6 +861,7 @@ def process(
         role_arn,
         tags,
         experiment_config=None,
+        pipeline_session=None,
     ):
         """Create an Amazon SageMaker processing job.
 
@@ -895,6 +901,7 @@ def process(
                 * If both `ExperimentName` and `TrialName` are not supplied the trial component
                 will be unassociated.
                 * `TrialComponentDisplayName` is used for display in Studio.
+            pipeline_session (`sagemaker.workflow.pipeline_session.PipelineSession`) A pipeline session
         """
         tags = _append_project_tags(tags)
         process_request = self._get_process_request(
@@ -910,6 +917,9 @@ def process(
             tags=tags,
             experiment_config=experiment_config,
         )
+        if pipeline_session:
+            pipeline_session.context = process_request
+            return
         LOGGER.info("Creating processing-job with name %s", job_name)
         LOGGER.debug("process request: %s", json.dumps(process_request, indent=4))
         self.sagemaker_client.create_processing_job(**process_request)
@@ -2045,6 +2055,7 @@ def create_tuning_job(
         training_config_list=None,
         warm_start_config=None,
         tags=None,
+        pipeline_session=None,
     ):
         """Create an Amazon SageMaker hyperparameter tuning job.
 
@@ -2081,7 +2092,9 @@ def create_tuning_job(
             warm_start_config=warm_start_config,
             tags=tags,
         )
-
+        if pipeline_session:
+            pipeline_session.context = tune_request
+            return
         LOGGER.info("Creating hyperparameter tuning job with name: %s", job_name)
         LOGGER.debug("tune request: %s", json.dumps(tune_request, indent=4))
         self.sagemaker_client.create_hyper_parameter_tuning_job(**tune_request)
@@ -2492,6 +2505,7 @@ def transform(
         tags,
         data_processing,
         model_client_config=None,
+        pipeline_session=None,
     ):
         """Create an Amazon SageMaker transform job.
 
@@ -2544,6 +2558,9 @@ def transform(
             data_processing=data_processing,
             model_client_config=model_client_config,
         )
+        if pipeline_session:
+            pipeline_session.context = transform_request
+            return
 
         LOGGER.info("Creating transform job with name: %s", job_name)
         LOGGER.debug("Transform request: %s", json.dumps(transform_request, indent=4))

src/sagemaker/spark/processing.py

@@ -249,7 +249,7 @@ def run(
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
-        super().run(
+        return super().run(
             submit_app,
             inputs,
             outputs,
@@ -868,7 +868,7 @@ def run(
             spark_event_logs_s3_uri=spark_event_logs_s3_uri,
         )
 
-        super().run(
+        return super().run(
             submit_app=submit_app,
             inputs=extended_inputs,
             outputs=extended_outputs,
@@ -1125,7 +1125,7 @@ def run(
             spark_event_logs_s3_uri=spark_event_logs_s3_uri,
         )
 
-        super().run(
+        return super().run(
             submit_app=submit_app,
             inputs=extended_inputs,
             outputs=extended_outputs,

src/sagemaker/transformer.py

@@ -17,6 +17,7 @@
 
 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.utils import base_name_from_image, name_from_base
 
 
@@ -106,6 +107,7 @@ def __init__(
 
         self.sagemaker_session = sagemaker_session or Session()
 
+    @runnable_by_pipeline
     def transform(
         self,
         data,
@@ -196,8 +198,10 @@ def transform(
         else:
             base_name = self.base_transform_job_name
 
-            if base_name is None:
+            if base_name is None and not is_pipeline_entities(self.model_name):
                 base_name = self._retrieve_base_name()
+            else:
+                base_name = f"transform-job"
 
             self._current_job_name = name_from_base(base_name)
 
@@ -370,6 +374,8 @@ def start_new(
             experiment_config,
             model_client_config,
         )
+        if type(transformer.sagemaker_session) is PipelineSession:
+            transform_args['pipeline_session'] = transformer.sagemaker_session
         transformer.sagemaker_session.transform(**transform_args)
 
         return cls(transformer.sagemaker_session, transformer._current_job_name)