fix: support idempotency for framework and spark processors (#3460)

Co-authored-by: Brock Wade <[email protected]>
Co-authored-by: Mufaddal Rohawala <[email protected]>

Author: 3 people

src/sagemaker/processing.py

@@ -23,6 +23,7 @@
 import logging
 from textwrap import dedent
 from typing import Dict, List, Optional, Union
+from copy import copy
 
 import attr
 
@@ -1830,14 +1831,17 @@ def _patch_inputs_with_payload(self, inputs, s3_payload) -> List[ProcessingInput
         #   a7399455f5386d83ddc5cb15c0db00c04bd518ec/src/sagemaker/processing.py#L425-L426
         if inputs is None:
             inputs = []
-        inputs.append(
+
+        # make a shallow copy of user inputs
+        patched_inputs = copy(inputs)
+        patched_inputs.append(
             ProcessingInput(
                 input_name="code",
                 source=s3_payload,
                 destination="/opt/ml/processing/input/code/",
             )
         )
-        return inputs
+        return patched_inputs
 
     def _set_entrypoint(self, command, user_script_name):
         """Framework processor override for setting processing job entrypoint.

src/sagemaker/spark/processing.py

@@ -30,6 +30,7 @@
 from enum import Enum
 from io import BytesIO
 from urllib.parse import urlparse
+from copy import copy
 
 from typing import Union, List, Dict, Optional
 
@@ -279,6 +280,10 @@ def run(
     def _extend_processing_args(self, inputs, outputs, **kwargs):
         """Extends processing job args such as inputs."""
 
+        # make a shallow copy of user outputs
+        outputs = outputs or []
+        extended_outputs = copy(outputs)
+
         if kwargs.get("spark_event_logs_s3_uri"):
             spark_event_logs_s3_uri = kwargs.get("spark_event_logs_s3_uri")
             self._validate_s3_uri(spark_event_logs_s3_uri)
@@ -297,16 +302,21 @@ def _extend_processing_args(self, inputs, outputs, **kwargs):
                 s3_upload_mode="Continuous",
             )
 
-            outputs = outputs or []
-            outputs.append(output)
+            extended_outputs.append(output)
+
+        # make a shallow copy of user inputs
+        inputs = inputs or []
+        extended_inputs = copy(inputs)
 
         if kwargs.get("configuration"):
             configuration = kwargs.get("configuration")
             self._validate_configuration(configuration)
-            inputs = inputs or []
-            inputs.append(self._stage_configuration(configuration))
+            extended_inputs.append(self._stage_configuration(configuration))
 
-        return inputs, outputs
+        return (
+            extended_inputs if extended_inputs else None,
+            extended_outputs if extended_outputs else None,
+        )
 
     def start_history_server(self, spark_event_logs_s3_uri=None):
         """Starts a Spark history server.
@@ -940,9 +950,16 @@ def _extend_processing_args(self, inputs, outputs, **kwargs):
             outputs: Processing outputs.
             kwargs: Additional keyword arguments passed to `super()`.
         """
+
+        if inputs is None:
+            inputs = []
+
+        # make a shallow copy of user inputs
+        extended_inputs = copy(inputs)
+
         self.command = [_SparkProcessorBase._default_command]
         extended_inputs = self._handle_script_dependencies(
-            inputs, kwargs.get("submit_py_files"), FileType.PYTHON
+            extended_inputs, kwargs.get("submit_py_files"), FileType.PYTHON
         )
         extended_inputs = self._handle_script_dependencies(
             extended_inputs, kwargs.get("submit_jars"), FileType.JAR
@@ -1199,8 +1216,14 @@ def _extend_processing_args(self, inputs, outputs, **kwargs):
         else:
             raise ValueError("submit_class is required")
 
+        if inputs is None:
+            inputs = []
+
+        # make a shallow copy of user inputs
+        extended_inputs = copy(inputs)
+
         extended_inputs = self._handle_script_dependencies(
-            inputs, kwargs.get("submit_jars"), FileType.JAR
+            extended_inputs, kwargs.get("submit_jars"), FileType.JAR
         )
         extended_inputs = self._handle_script_dependencies(
             extended_inputs, kwargs.get("submit_files"), FileType.FILE

src/sagemaker/workflow/utilities.py

@@ -114,11 +114,12 @@ def get_code_hash(step: Entity) -> str:
     if isinstance(step, ProcessingStep) and step.step_args:
         kwargs = step.step_args.func_kwargs
         source_dir = kwargs.get("source_dir")
+        submit_class = kwargs.get("submit_class")
         dependencies = get_processing_dependencies(
             [
                 kwargs.get("dependencies"),
                 kwargs.get("submit_py_files"),
-                kwargs.get("submit_class"),
+                [submit_class] if submit_class else None,
                 kwargs.get("submit_jars"),
                 kwargs.get("submit_files"),
             ]
@@ -168,7 +169,7 @@ def get_processing_code_hash(code: str, source_dir: str, dependencies: List[str]
         str: A hash string representing the unique code artifact(s) for the step
     """
 
-    # FrameworkProcessor
+    # If FrameworkProcessor contains source_dir
     if source_dir:
         source_dir_url = urlparse(source_dir)
         if source_dir_url.scheme == "" or source_dir_url.scheme == "file":
@@ -400,5 +401,5 @@ def execute_job_functions(step_args: _StepArguments):
     """
 
     chained_args = step_args.func(*step_args.func_args, **step_args.func_kwargs)
-    if chained_args:
+    if isinstance(chained_args, _StepArguments):
         execute_job_functions(chained_args)

tests/data/spark/code/java/TestJarFile.jar

@@ -17,7 +17,7 @@
 
 import pytest
 
-from mock import Mock
+from mock import Mock, patch
 
 from sagemaker import s3
 from sagemaker.workflow.condition_step import ConditionStep
@@ -78,6 +78,7 @@ def test_pipeline_create_with_parallelism_config(sagemaker_session_mock, role_ar
     )
 
 
+@patch("sagemaker.s3.S3Uploader.upload_string_as_file_body")
 def test_large_pipeline_create(sagemaker_session_mock, role_arn):
     parameter = ParameterString("MyStr")
     pipeline = Pipeline(
@@ -87,8 +88,6 @@ def test_large_pipeline_create(sagemaker_session_mock, role_arn):
         sagemaker_session=sagemaker_session_mock,
     )
 
-    s3.S3Uploader.upload_string_as_file_body = Mock()
-
     pipeline.create(role_arn=role_arn)
 
     assert s3.S3Uploader.upload_string_as_file_body.called_with(
@@ -151,6 +150,7 @@ def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_ar
     )
 
 
+@patch("sagemaker.s3.S3Uploader.upload_string_as_file_body")
 def test_large_pipeline_update(sagemaker_session_mock, role_arn):
     parameter = ParameterString("MyStr")
     pipeline = Pipeline(
@@ -160,8 +160,6 @@ def test_large_pipeline_update(sagemaker_session_mock, role_arn):
         sagemaker_session=sagemaker_session_mock,
     )
 
-    s3.S3Uploader.upload_string_as_file_body = Mock()
-
     pipeline.create(role_arn=role_arn)
 
     assert s3.S3Uploader.upload_string_as_file_body.called_with(