feat: Support selective pipeline execution between function step and regular step

Author: qidewenwhen

src/sagemaker/workflow/function_step.py

@@ -33,12 +33,11 @@
     PipelineVariable,
 )
 
-from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.retry import RetryPolicy
 from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
 from sagemaker.workflow.step_collections import StepCollection
-from sagemaker.workflow.step_outputs import StepOutput
+from sagemaker.workflow.step_outputs import StepOutput, get_step
 from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context
 
 from sagemaker.s3_utils import s3_path_join
@@ -275,16 +274,11 @@ def expr(self) -> RequestType:
 
     def _to_json_get(self) -> JsonGet:
         """Expression structure for workflow service calls using JsonGet resolution."""
-        from sagemaker.remote_function.core.stored_function import (
-            JSON_SERIALIZED_RESULT_KEY,
-            RESULTS_FOLDER,
-            JSON_RESULTS_FILE,
-        )
+        from sagemaker.remote_function.core.stored_function import JSON_SERIALIZED_RESULT_KEY
 
         if not self._step.name:
             raise ValueError("Step name is not defined.")
 
-        s3_root_uri = self._step._job_settings.s3_root_uri
         # Resolve json path --
         #   Deserializer will be able to resolve a JsonGet using path "Return[1]" to
         #   access value 10 from following serialized JSON:
@@ -308,15 +302,8 @@ def _to_json_get(self) -> JsonGet:
 
         return JsonGet(
             s3_uri=Join(
-                "/",
-                [
-                    s3_root_uri,
-                    ExecutionVariables.PIPELINE_NAME,
-                    ExecutionVariables.PIPELINE_EXECUTION_ID,
-                    self._step.name,
-                    RESULTS_FOLDER,
-                    JSON_RESULTS_FILE,
-                ],
+                on="",
+                values=[get_step(self)._properties.OutputDataConfig.S3OutputPath],
             ),
             json_path=_resolved_reference_path,
             step=self._step,

src/sagemaker/workflow/functions.py

@@ -21,7 +21,7 @@
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.execution_variables import ExecutionVariable
 from sagemaker.workflow.parameters import Parameter
-from sagemaker.workflow.properties import PropertyFile
+from sagemaker.workflow.properties import PropertyFile, Properties
 
 if TYPE_CHECKING:
     from sagemaker.workflow.steps import Step
@@ -172,9 +172,9 @@ def _validate_json_get_s3_uri(self):
         for join_arg in s3_uri.values:
             if not is_pipeline_variable(join_arg):
                 continue
-            if not isinstance(join_arg, (Parameter, ExecutionVariable)):
+            if not isinstance(join_arg, (Parameter, ExecutionVariable, Properties)):
                 raise ValueError(
                     f"Invalid JsonGet function {self.expr}. "
-                    f"The Join values in JsonGet's s3_uri can only be a primitive object "
-                    f"or Parameter or ExecutionVariable."
+                    f"The Join values in JsonGet's s3_uri can only be a primitive object, "
+                    f"Parameter, ExecutionVariable or Properties."
                 )

tests/integ/sagemaker/workflow/test_selective_execution.py

@@ -16,6 +16,7 @@
 
 import pytest
 
+from sagemaker.processing import ProcessingInput
 from tests.integ import DATA_DIR
 from sagemaker.sklearn import SKLearnProcessor
 from sagemaker.workflow.step_outputs import get_step
@@ -199,3 +200,102 @@ def func_2(arg):
             pipeline.delete()
         except Exception:
             pass
+
+
+def test_decorator_step_data_referenced_by_other_steps(
+    pipeline_session,
+    role,
+    pipeline_name,
+    region_name,
+    dummy_container_without_error,
+    sklearn_latest_version,
+):
+    # Test Selective Pipeline Execution on function step -> [select: regular step]
+    os.environ["AWS_DEFAULT_REGION"] = region_name
+    processing_job_instance_counts = 2
+
+    @step(
+        name="step1",
+        role=role,
+        image_uri=dummy_container_without_error,
+        instance_type=INSTANCE_TYPE,
+        keep_alive_period_in_seconds=60,
+    )
+    def func(var: int):
+        return 1, var
+
+    step_output = func(processing_job_instance_counts)
+
+    script_path = os.path.join(DATA_DIR, "dummy_script.py")
+    input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
+    inputs = [
+        ProcessingInput(source=input_file_path, destination="/opt/ml/processing/inputs/"),
+    ]
+
+    sklearn_processor = SKLearnProcessor(
+        framework_version=sklearn_latest_version,
+        role=role,
+        instance_type=INSTANCE_TYPE,
+        instance_count=step_output[1],
+        command=["python3"],
+        sagemaker_session=pipeline_session,
+        base_job_name="test-sklearn",
+    )
+
+    step_args = sklearn_processor.run(
+        inputs=inputs,
+        code=script_path,
+    )
+    process_step = ProcessingStep(
+        name="MyProcessStep",
+        step_args=step_args,
+    )
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        steps=[process_step],
+        sagemaker_session=pipeline_session,
+    )
+
+    try:
+        execution, _ = create_and_execute_pipeline(
+            pipeline=pipeline,
+            pipeline_name=pipeline_name,
+            region_name=region_name,
+            role=role,
+            no_of_steps=2,
+            last_step_name=process_step.name,
+            execution_parameters=dict(),
+            step_status="Succeeded",
+            wait_duration=1000,  # seconds
+        )
+
+        _, execution_steps2 = create_and_execute_pipeline(
+            pipeline=pipeline,
+            pipeline_name=pipeline_name,
+            region_name=region_name,
+            role=role,
+            no_of_steps=2,
+            last_step_name=process_step.name,
+            execution_parameters=dict(),
+            step_status="Succeeded",
+            wait_duration=1000,  # seconds
+            selective_execution_config=SelectiveExecutionConfig(
+                source_pipeline_execution_arn=execution.arn,
+                selected_steps=[process_step.name],
+            ),
+        )
+
+        execution_proc_job = pipeline_session.describe_processing_job(
+            execution_steps2[0]["Metadata"]["ProcessingJob"]["Arn"].split("/")[-1]
+        )
+        assert (
+            execution_proc_job["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
+            == processing_job_instance_counts
+        )
+
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass

tests/unit/sagemaker/workflow/test_function_step.py

@@ -328,15 +328,8 @@ def func():
             "Path": "Result['some_key'][0]",
             "S3Uri": {
                 "Std:Join": {
-                    "On": "/",
-                    "Values": [
-                        "s3://bucket",
-                        {"Get": "Execution.PipelineName"},
-                        {"Get": "Execution.PipelineExecutionId"},
-                        "step_name",
-                        "results",
-                        "results.json",
-                    ],
+                    "On": "",
+                    "Values": [{"Get": "Steps.step_name.OutputDataConfig.S3OutputPath"}],
                 }
             },
         }
@@ -352,15 +345,8 @@ def func():
             "Path": "Result[0]['some_key']",
             "S3Uri": {
                 "Std:Join": {
-                    "On": "/",
-                    "Values": [
-                        "s3://bucket",
-                        {"Get": "Execution.PipelineName"},
-                        {"Get": "Execution.PipelineExecutionId"},
-                        "step_name",
-                        "results",
-                        "results.json",
-                    ],
+                    "On": "",
+                    "Values": [{"Get": "Steps.step_name.OutputDataConfig.S3OutputPath"}],
                 }
             },
         }

tests/unit/sagemaker/workflow/test_functions.py

@@ -279,8 +279,8 @@ def test_json_get_invalid_s3_uri_not_join():
 
 def test_json_get_invalid_s3_uri_with_invalid_pipeline_variable(sagemaker_session):
     with pytest.raises(ValueError) as e:
-        JsonGet(s3_uri=Join(on="/", values=["s3:/", Properties(step_name="test")]))
+        JsonGet(s3_uri=Join(on="/", values=["s3:/", Join()]))
     assert (
-        "The Join values in JsonGet's s3_uri can only be a primitive object or Parameter or ExecutionVariable."
-        in str(e.value)
+        "The Join values in JsonGet's s3_uri can only be a primitive object, "
+        "Parameter, ExecutionVariable or Properties." in str(e.value)
     )