feat: Support selective pipeline execution for function step

Author: qidewenwhen

src/sagemaker/remote_function/core/pipeline_variables.py

@@ -19,6 +19,7 @@
 
 from sagemaker.s3 import s3_path_join
 from sagemaker.remote_function.core.serialization import deserialize_obj_from_s3
+from sagemaker.workflow.step_outputs import get_step
 
 
 @dataclass
@@ -166,6 +167,7 @@ def __init__(
         hmac_key: str,
         parameter_resolver: _ParameterResolver,
         execution_variable_resolver: _ExecutionVariableResolver,
+        properties_resolver: _PropertiesResolver,
         s3_base_uri: str,
         **settings,
     ):
@@ -184,6 +186,7 @@ def __init__(
         self._s3_base_uri = s3_base_uri
         self._parameter_resolver = parameter_resolver
         self._execution_variable_resolver = execution_variable_resolver
+        self._properties_resolver = properties_resolver
         # different delayed returns can have the same uri, so we need to dedupe
         uris = {
             self._resolve_delayed_return_uri(delayed_return) for delayed_return in delayed_returns
@@ -223,6 +226,8 @@ def _resolve_delayed_return_uri(self, delayed_return: _DelayedReturn):
                 uri.append(self._execution_variable_resolver.resolve(component))
             elif isinstance(component, _S3BaseUriIdentifier):
                 uri.append(self._s3_base_uri)
+            elif isinstance(component, _Properties):
+                uri.append(self._properties_resolver.resolve(component))
             else:
                 uri.append(component)
         return s3_path_join(*uri)
@@ -276,6 +281,7 @@ def resolve_pipeline_variables(
         hmac_key=hmac_key,
         parameter_resolver=parameter_resolver,
         execution_variable_resolver=execution_variable_resolver,
+        properties_resolver=properties_resolver,
         s3_base_uri=s3_base_uri,
         **settings,
     )
@@ -322,39 +328,43 @@ def convert_pipeline_variables_to_pickleable(func_args: Tuple, func_kwargs: Dict
         func_args: function args.
         func_kwargs: function kwargs.
     """
+    converted_func_args = tuple(_convert_pipeline_variable_to_pickleable(arg) for arg in func_args)
+    converted_func_kwargs = {
+        key: _convert_pipeline_variable_to_pickleable(arg) for key, arg in func_kwargs.items()
+    }
 
-    from sagemaker.workflow.entities import PipelineVariable
-
-    from sagemaker.workflow.execution_variables import ExecutionVariables
-
-    from sagemaker.workflow.function_step import DelayedReturn
-
-    # Notes:
-    # 1. The s3_base_uri = s3_root_uri + pipeline_name, but the two may be unknown
-    # when defining function steps. After step-level arg serialization,
-    # it's hard to update the s3_base_uri in pipeline compile time.
-    # Thus set a placeholder: _S3BaseUriIdentifier, and let the runtime job to resolve it.
-    # 2. For saying s3_root_uri is unknown, it's because when defining function steps,
-    # the pipeline's sagemaker_session is not passed in, but the default s3_root_uri
-    # should be retrieved from the pipeline's sagemaker_session.
-    def convert(arg):
-        if isinstance(arg, DelayedReturn):
-            return _DelayedReturn(
-                uri=[
-                    _S3BaseUriIdentifier(),
-                    ExecutionVariables.PIPELINE_EXECUTION_ID._pickleable,
-                    arg._step.name,
-                    "results",
-                ],
-                reference_path=arg._reference_path,
-            )
+    return converted_func_args, converted_func_kwargs
 
-        if isinstance(arg, PipelineVariable):
-            return arg._pickleable
 
-        return arg
+def _convert_pipeline_variable_to_pickleable(arg):
+    """Convert a pipeline variable to pickleable."""
+    from sagemaker.workflow.entities import PipelineVariable
 
-    converted_func_args = tuple(convert(arg) for arg in func_args)
-    converted_func_kwargs = {key: convert(arg) for key, arg in func_kwargs.items()}
+    from sagemaker.workflow.function_step import DelayedReturn
 
-    return converted_func_args, converted_func_kwargs
+    if isinstance(arg, DelayedReturn):
+        # Notes:
+        # 1. The s3_base_uri = s3_root_uri + pipeline_name, but the two may be unknown
+        # when defining function steps. After step-level arg serialization,
+        # it's hard to update the s3_base_uri in pipeline compile time.
+        # Thus set a placeholder: _S3BaseUriIdentifier, and let the runtime job to resolve it.
+        # 2. For saying s3_root_uri is unknown, it's because when defining function steps,
+        # the pipeline's sagemaker_session is not passed in, but the default s3_root_uri
+        # should be retrieved from the pipeline's sagemaker_session.
+
+        container_args = get_step(arg)._properties.AlgorithmSpecification.ContainerArguments
+        execution_id = container_args[11]._pickleable
+        return _DelayedReturn(
+            uri=[
+                _S3BaseUriIdentifier(),
+                execution_id,
+                arg._step.name,
+                "results",
+            ],
+            reference_path=arg._reference_path,
+        )
+
+    if isinstance(arg, PipelineVariable):
+        return arg._pickleable
+
+    return arg

src/sagemaker/remote_function/job.py

@@ -59,7 +59,11 @@
 from sagemaker.s3 import s3_path_join, S3Uploader
 from sagemaker import vpc_utils
 from sagemaker.remote_function.core.stored_function import StoredFunction, _SerializedData
-from sagemaker.remote_function.core.pipeline_variables import Context
+from sagemaker.remote_function.core.pipeline_variables import (
+    Context,
+    _convert_pipeline_variable_to_pickleable,
+    _Properties,
+)
 from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
     RuntimeEnvironmentManager,
     _DependencySettings,
@@ -72,6 +76,7 @@
     copy_workdir,
     resolve_custom_file_filter_from_config_file,
 )
+from sagemaker.workflow.function_step import DelayedReturn
 
 if TYPE_CHECKING:
     from sagemaker.workflow.entities import PipelineVariable
@@ -804,6 +809,13 @@ def compile(
                 if isinstance(arg, (Parameter, ExecutionVariable, Properties)):
                     container_args.extend([arg.expr["Get"], arg.to_string()])
 
+                if isinstance(arg, DelayedReturn):
+                    uri = _convert_pipeline_variable_to_pickleable(arg).uri
+                    for uri_element in uri:
+                        if not isinstance(uri_element, _Properties):
+                            continue
+                        container_args.extend([uri_element.path, {"Get": uri_element.path}])
+
         if run_info is not None:
             container_args.extend(["--run_in_context", json.dumps(dataclasses.asdict(run_info))])
         elif _RunContext.get_current_run() is not None:

src/sagemaker/workflow/function_step.py

@@ -34,6 +34,7 @@
 )
 
 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
@@ -100,6 +101,9 @@ def __init__(
         self._step_kwargs = kwargs
 
         self.__job_settings = None
+        self._properties = Properties(
+            step_name=name, step=self, shape_name="DescribeTrainingJobResponse"
+        )
 
         (
             self._converted_func_args,

tests/integ/sagemaker/workflow/helpers.py

@@ -39,18 +39,24 @@ def create_and_execute_pipeline(
     step_result_type=None,
     step_result_value=None,
     wait_duration=400,  # seconds
+    selective_execution_config=None,
 ):
-    response = pipeline.create(role)
-
-    create_arn = response["PipelineArn"]
-    assert re.match(
-        rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
-        create_arn,
+    create_arn = None
+    if not selective_execution_config:
+        response = pipeline.create(role)
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+    execution = pipeline.start(
+        parameters=execution_parameters, selective_execution_config=selective_execution_config
     )
 
-    execution = pipeline.start(parameters=execution_parameters)
-    response = execution.describe()
-    assert response["PipelineArn"] == create_arn
+    if create_arn:
+        response = execution.describe()
+        assert response["PipelineArn"] == create_arn
 
     wait_pipeline_execution(execution=execution, delay=20, max_attempts=int(wait_duration / 20))
 
@@ -71,6 +77,16 @@ def create_and_execute_pipeline(
     if step_result_value:
         result = execution.result(execution_steps[0]["StepName"])
         assert result == step_result_value, f"Expected {step_result_value}, instead found {result}"
+
+    if selective_execution_config:
+        for exe_step in execution_steps:
+            if exe_step["StepName"] in selective_execution_config.selected_steps:
+                continue
+            assert (
+                exe_step["SelectiveExecutionResult"]["SourcePipelineExecutionArn"]
+                == selective_execution_config.source_pipeline_execution_arn
+            )
+
     return execution, execution_steps
 
 

tests/integ/sagemaker/workflow/test_step_decorator.py

@@ -39,6 +39,7 @@
     ParameterFloat,
     ParameterBoolean,
 )
+from sagemaker.workflow.selective_execution_config import SelectiveExecutionConfig
 from sagemaker.workflow.step_outputs import get_step
 from sagemaker.workflow.steps import ProcessingStep
 
@@ -246,6 +247,20 @@ def sum(a, b):
     )
 
     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="sum",
+            execution_parameters=dict(),
+            step_status="Succeeded",
+            step_result_type=int,
+            step_result_value=7,
+        )
+
+        # Test Selective Pipeline Execution on function step1 -> [select: function step2]
         create_and_execute_pipeline(
             pipeline=pipeline,
             pipeline_name=pipeline_name,
@@ -256,7 +271,13 @@ def sum(a, b):
             execution_parameters=dict(),
             step_status="Succeeded",
             step_result_type=int,
+            step_result_value=7,
+            selective_execution_config=SelectiveExecutionConfig(
+                source_pipeline_execution_arn=execution.arn,
+                selected_steps=[get_step(step_output_b).name],
+            ),
         )
+
     finally:
         try:
             pipeline.delete()
@@ -379,7 +400,7 @@ def func_2(*args):
     )
 
     try:
-        create_and_execute_pipeline(
+        execution, _ = create_and_execute_pipeline(
             pipeline=pipeline,
             pipeline_name=pipeline_name,
             region_name=region_name,
@@ -392,6 +413,24 @@ def func_2(*args):
             step_result_value=(3, True, 2.0, "string", "Completed", 3),
             wait_duration=600,
         )
+        # Test Selective Pipeline Execution on regular step -> [select: function step]
+        execution, _ = create_and_execute_pipeline(
+            pipeline=pipeline,
+            pipeline_name=pipeline_name,
+            region_name=region_name,
+            role=role,
+            no_of_steps=3,
+            last_step_name="func",
+            execution_parameters=dict(param_a=10),
+            step_status="Succeeded",
+            step_result_type=tuple,
+            step_result_value=(10, True, 2.0, "string", "Completed", 3),
+            wait_duration=600,
+            selective_execution_config=SelectiveExecutionConfig(
+                source_pipeline_execution_arn=execution.arn,
+                selected_steps=[get_step(final_output).name],
+            ),
+        )
 
     finally:
         try: