feature: Caching Improvements for SM Pipeline Workflows (#3441)

Author: brockwade633

src/sagemaker/algorithm.py

@@ -27,6 +27,7 @@
 from sagemaker.predictor import Predictor
 from sagemaker.session import Session
 from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 
 from sagemaker.workflow import is_pipeline_variable
 
@@ -429,6 +430,7 @@ def _prepare_for_training(self, job_name=None):
 
         super(AlgorithmEstimator, self)._prepare_for_training(job_name)
 
+    @runnable_by_pipeline
     def fit(
         self,
         inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None,

src/sagemaker/estimator.py

@@ -778,13 +778,11 @@ def _stage_user_code_in_s3(self) -> str:
         if is_pipeline_variable(self.output_path):
             if self.code_location is None:
                 code_bucket = self.sagemaker_session.default_bucket()
-                code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                code_s3_prefix = self._assign_s3_prefix()
                 kms_key = None
             else:
                 code_bucket, key_prefix = parse_s3_url(self.code_location)
-                code_s3_prefix = "/".join(
-                    filter(None, [key_prefix, self._current_job_name, "source"])
-                )
+                code_s3_prefix = self._assign_s3_prefix(key_prefix)
 
                 output_bucket = self.sagemaker_session.default_bucket()
                 kms_key = self.output_kms_key if code_bucket == output_bucket else None
@@ -793,24 +791,20 @@ def _stage_user_code_in_s3(self) -> str:
             if local_mode:
                 if self.code_location is None:
                     code_bucket = self.sagemaker_session.default_bucket()
-                    code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                    code_s3_prefix = self._assign_s3_prefix()
                     kms_key = None
                 else:
                     code_bucket, key_prefix = parse_s3_url(self.code_location)
-                    code_s3_prefix = "/".join(
-                        filter(None, [key_prefix, self._current_job_name, "source"])
-                    )
+                    code_s3_prefix = self._assign_s3_prefix(key_prefix)
                     kms_key = None
             else:
                 if self.code_location is None:
                     code_bucket, _ = parse_s3_url(self.output_path)
-                    code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                    code_s3_prefix = self._assign_s3_prefix()
                     kms_key = self.output_kms_key
                 else:
                     code_bucket, key_prefix = parse_s3_url(self.code_location)
-                    code_s3_prefix = "/".join(
-                        filter(None, [key_prefix, self._current_job_name, "source"])
-                    )
+                    code_s3_prefix = self._assign_s3_prefix(key_prefix)
 
                     output_bucket, _ = parse_s3_url(self.output_path)
                     kms_key = self.output_kms_key if code_bucket == output_bucket else None
@@ -827,6 +821,36 @@ def _stage_user_code_in_s3(self) -> str:
             settings=self.sagemaker_session.settings,
         )
 
+    def _assign_s3_prefix(self, key_prefix=""):
+        """Include pipeline name+step name instead of job name in s3 path
+
+        Assign new s3 path structure if within a pipeline workflow that has
+            set the _pipeline_config and respective name/hash variables
+
+        Args:
+            key_prefix (str): Prefix for the S3 key, often netloc of url:
+            https://docs.python.org/3.9/library/urllib.parse.html#urllib.parse.netloc
+
+        Returns:
+            str: S3 path prefix that occurs before filename
+        """
+        from sagemaker.workflow.utilities import _pipeline_config
+
+        code_s3_prefix = "/".join(filter(None, [key_prefix, self._current_job_name, "source"]))
+        if _pipeline_config and _pipeline_config.code_hash:
+            code_s3_prefix = "/".join(
+                filter(
+                    None,
+                    [
+                        key_prefix,
+                        _pipeline_config.pipeline_name,
+                        "code",
+                        _pipeline_config.code_hash,
+                    ],
+                )
+            )
+        return code_s3_prefix
+
     def _prepare_rules(self):
         """Rules list includes both debugger and profiler rules.
 
@@ -1539,7 +1563,6 @@ def model_data(self):
             model_uri = os.path.join(
                 self.output_path, self._current_job_name, "output", "model.tar.gz"
             )
-
         return model_uri
 
     @abstractmethod

src/sagemaker/processing.py

@@ -35,7 +35,9 @@
 from sagemaker.utils import base_name_from_image, get_config_value, name_from_base
 from sagemaker.session import Session
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.workflow.functions import Join
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
+from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.dataset_definition.inputs import S3Input, DatasetDefinition
 from sagemaker.apiutils._base_types import ApiObject
@@ -314,6 +316,8 @@ def _normalize_inputs(self, inputs=None, kms_key=None):
         Raises:
             TypeError: if the inputs are not ``ProcessingInput`` objects.
         """
+        from sagemaker.workflow.utilities import _pipeline_config
+
         # Initialize a list of normalized ProcessingInput objects.
         normalized_inputs = []
         if inputs is not None:
@@ -335,13 +339,23 @@ def _normalize_inputs(self, inputs=None, kms_key=None):
                 # and save the S3 uri in the ProcessingInput source.
                 parse_result = urlparse(file_input.s3_input.s3_uri)
                 if parse_result.scheme != "s3":
-                    desired_s3_uri = s3.s3_path_join(
-                        "s3://",
-                        self.sagemaker_session.default_bucket(),
-                        self._current_job_name,
-                        "input",
-                        file_input.input_name,
-                    )
+                    if _pipeline_config:
+                        desired_s3_uri = s3.s3_path_join(
+                            "s3://",
+                            self.sagemaker_session.default_bucket(),
+                            _pipeline_config.pipeline_name,
+                            _pipeline_config.step_name,
+                            "input",
+                            file_input.input_name,
+                        )
+                    else:
+                        desired_s3_uri = s3.s3_path_join(
+                            "s3://",
+                            self.sagemaker_session.default_bucket(),
+                            self._current_job_name,
+                            "input",
+                            file_input.input_name,
+                        )
                     s3_uri = s3.S3Uploader.upload(
                         local_path=file_input.s3_input.s3_uri,
                         desired_s3_uri=desired_s3_uri,
@@ -369,6 +383,8 @@ def _normalize_outputs(self, outputs=None):
             TypeError: if the outputs are not ``ProcessingOutput`` objects.
         """
         # Initialize a list of normalized ProcessingOutput objects.
+        from sagemaker.workflow.utilities import _pipeline_config
+
         normalized_outputs = []
         if outputs is not None:
             # Iterate through the provided list of outputs.
@@ -384,13 +400,27 @@ def _normalize_outputs(self, outputs=None):
                 # If the output's destination is not an s3_uri, create one.
                 parse_result = urlparse(output.destination)
                 if parse_result.scheme != "s3":
-                    s3_uri = s3.s3_path_join(
-                        "s3://",
-                        self.sagemaker_session.default_bucket(),
-                        self._current_job_name,
-                        "output",
-                        output.output_name,
-                    )
+                    if _pipeline_config:
+                        s3_uri = Join(
+                            on="/",
+                            values=[
+                                "s3:/",
+                                self.sagemaker_session.default_bucket(),
+                                _pipeline_config.pipeline_name,
+                                ExecutionVariables.PIPELINE_EXECUTION_ID,
+                                _pipeline_config.step_name,
+                                "output",
+                                output.output_name,
+                            ],
+                        )
+                    else:
+                        s3_uri = s3.s3_path_join(
+                            "s3://",
+                            self.sagemaker_session.default_bucket(),
+                            self._current_job_name,
+                            "output",
+                            output.output_name,
+                        )
                     output.destination = s3_uri
                 normalized_outputs.append(output)
         return normalized_outputs
@@ -507,6 +537,11 @@ def get_run_args(
             arguments (list[str]): A list of string arguments to be passed to a
                 processing job (default: None).
         """
+        logger.warning(
+            "This function has been deprecated and could break pipeline step caching. "
+            "We recommend using the run() function directly with pipeline sessions"
+            "to access step arguments."
+        )
         return RunArgs(code=code, inputs=inputs, outputs=outputs, arguments=arguments)
 
     @runnable_by_pipeline
@@ -679,13 +714,24 @@ def _upload_code(self, code, kms_key=None):
             str: The S3 URI of the uploaded file or directory.
 
         """
-        desired_s3_uri = s3.s3_path_join(
-            "s3://",
-            self.sagemaker_session.default_bucket(),
-            self._current_job_name,
-            "input",
-            self._CODE_CONTAINER_INPUT_NAME,
-        )
+        from sagemaker.workflow.utilities import _pipeline_config
+
+        if _pipeline_config and _pipeline_config.code_hash:
+            desired_s3_uri = s3.s3_path_join(
+                "s3://",
+                self.sagemaker_session.default_bucket(),
+                _pipeline_config.pipeline_name,
+                self._CODE_CONTAINER_INPUT_NAME,
+                _pipeline_config.code_hash,
+            )
+        else:
+            desired_s3_uri = s3.s3_path_join(
+                "s3://",
+                self.sagemaker_session.default_bucket(),
+                self._current_job_name,
+                "input",
+                self._CODE_CONTAINER_INPUT_NAME,
+            )
         return s3.S3Uploader.upload(
             local_path=code,
             desired_s3_uri=desired_s3_uri,
@@ -1499,6 +1545,12 @@ def get_run_args(
             job_name (str): Processing job name. If not specified, the processor generates
                 a default job name, based on the base job name and current timestamp.
         """
+        logger.warning(
+            "This function has been deprecated and could break pipeline step caching. "
+            "We recommend using the run() function directly with pipeline sessions"
+            "to access step arguments."
+        )
+
         # When job_name is None, the job_name to upload code (+payload) will
         # differ from job_name used by run().
         s3_runproc_sh, inputs, job_name = self._pack_and_upload_code(
@@ -1512,6 +1564,7 @@ def get_run_args(
             arguments=arguments,
         )
 
+    @runnable_by_pipeline
     def run(  # type: ignore[override]
         self,
         code: str,

src/sagemaker/spark/processing.py

@@ -42,6 +42,7 @@
 from sagemaker.spark import defaults
 
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.functions import Join
 
@@ -211,6 +212,7 @@ def get_run_args(
             arguments=arguments,
         )
 
+    @runnable_by_pipeline
     def run(
         self,
         submit_app,
@@ -399,12 +401,22 @@ def _stage_configuration(self, configuration):
         Args:
             configuration (Dict): the configuration dict for the EMR application configuration.
         """
+        from sagemaker.workflow.utilities import _pipeline_config
 
         serialized_configuration = BytesIO(json.dumps(configuration).encode("utf-8"))
-        s3_uri = (
-            f"s3://{self.sagemaker_session.default_bucket()}/{self._current_job_name}/"
-            f"input/{self._conf_container_input_name}/{self._conf_file_name}"
-        )
+
+        if _pipeline_config and _pipeline_config.config_hash:
+            s3_uri = (
+                f"s3://{self.sagemaker_session.default_bucket()}/{_pipeline_config.pipeline_name}/"
+                f"{_pipeline_config.step_name}/input/"
+                f"{self._conf_container_input_name}/{_pipeline_config.config_hash}/"
+                f"{self._conf_file_name}"
+            )
+        else:
+            s3_uri = (
+                f"s3://{self.sagemaker_session.default_bucket()}/{self._current_job_name}/"
+                f"input/{self._conf_container_input_name}/{self._conf_file_name}"
+            )
 
         S3Uploader.upload_string_as_file_body(
             body=serialized_configuration,
@@ -443,11 +455,6 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
         if not input_channel_name:
             raise ValueError("input_channel_name value may not be empty.")
 
-        input_channel_s3_uri = (
-            f"s3://{self.sagemaker_session.default_bucket()}"
-            f"/{self._current_job_name}/input/{input_channel_name}"
-        )
-
         use_input_channel = False
         spark_opt_s3_uris = []
         spark_opt_s3_uris_has_pipeline_var = False
@@ -481,6 +488,19 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
 
             # If any local files were found and copied, upload the temp directory to S3
             if os.listdir(tmpdir):
+                from sagemaker.workflow.utilities import _pipeline_config
+
+                if _pipeline_config and _pipeline_config.code_hash:
+                    input_channel_s3_uri = (
+                        f"s3://{self.sagemaker_session.default_bucket()}"
+                        f"/{_pipeline_config.pipeline_name}/code/{_pipeline_config.code_hash}"
+                        f"/{input_channel_name}"
+                    )
+                else:
+                    input_channel_s3_uri = (
+                        f"s3://{self.sagemaker_session.default_bucket()}"
+                        f"/{self._current_job_name}/input/{input_channel_name}"
+                    )
                 logger.info(
                     "Uploading dependencies from tmpdir %s to S3 %s", tmpdir, input_channel_s3_uri
                 )
@@ -824,6 +844,7 @@ def get_run_args(
             arguments=arguments,
         )
 
+    @runnable_by_pipeline
     def run(
         self,
         submit_app: str,
@@ -1083,6 +1104,7 @@ def get_run_args(
             arguments=arguments,
         )
 
+    @runnable_by_pipeline
     def run(
         self,
         submit_app: str,