feat: Support selective pipeline execution for function step

Author: qidewenwhen

src/sagemaker/local/entities.py

@@ -765,7 +765,6 @@ def result(self, step_name: str):
         return get_function_step_result(
             step_name=step_name,
             step_list=self.list_steps()["PipelineExecutionSteps"],
-            execution_id=self.pipeline_execution_name,
             sagemaker_session=self.local_session,
         )
 

src/sagemaker/remote_function/core/pipeline_variables.py

@@ -15,10 +15,10 @@
 
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
-from typing import Any, Union, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple
 
-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
@@ -77,22 +77,11 @@ class _ExecutionVariable:
     name: str
 
 
-@dataclass
-class _S3BaseUriIdentifier:
-    """Identifies that the class refers to function step s3 base uri.
-
-    The s3_base_uri = s3_root_uri + pipeline_name.
-    This identifier is resolved in function step runtime by SDK.
-    """
-
-    NAME = "S3_BASE_URI"
-
-
 @dataclass
 class _DelayedReturn:
     """Delayed return from a function."""
 
-    uri: List[Union[str, _Parameter, _ExecutionVariable]]
+    uri: _Properties
     reference_path: Tuple = field(default_factory=tuple)
 
 
@@ -164,26 +153,18 @@ def __init__(
         self,
         delayed_returns: List[_DelayedReturn],
         hmac_key: str,
-        parameter_resolver: _ParameterResolver,
-        execution_variable_resolver: _ExecutionVariableResolver,
-        s3_base_uri: str,
+        properties_resolver: _PropertiesResolver,
         **settings,
     ):
         """Resolve delayed return.
 
         Args:
             delayed_returns: list of delayed returns to resolve.
             hmac_key: key used to encrypt serialized and deserialized function and arguments.
-            parameter_resolver: resolver used to pipeline parameters.
-            execution_variable_resolver: resolver used to resolve execution variables.
-            s3_base_uri (str): the s3 base uri of the function step that
-                the serialized artifacts will be uploaded to.
-                The s3_base_uri = s3_root_uri + pipeline_name.
+            properties_resolver: resolver used to resolve step properties.
             **settings: settings to pass to the deserialization function.
         """
-        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
@@ -214,18 +195,7 @@ def resolve(self, delayed_return: _DelayedReturn) -> Any:
 
     def _resolve_delayed_return_uri(self, delayed_return: _DelayedReturn):
         """Resolve the s3 uri of the delayed return."""
-
-        uri = []
-        for component in delayed_return.uri:
-            if isinstance(component, _Parameter):
-                uri.append(self._parameter_resolver.resolve(component))
-            elif isinstance(component, _ExecutionVariable):
-                uri.append(self._execution_variable_resolver.resolve(component))
-            elif isinstance(component, _S3BaseUriIdentifier):
-                uri.append(self._s3_base_uri)
-            else:
-                uri.append(component)
-        return s3_path_join(*uri)
+        return self._properties_resolver.resolve(delayed_return.uri)
 
 
 def _retrieve_child_item(delayed_return: _DelayedReturn, deserialized_obj: Any):
@@ -241,7 +211,6 @@ def resolve_pipeline_variables(
     func_args: Tuple,
     func_kwargs: Dict,
     hmac_key: str,
-    s3_base_uri: str,
     **settings,
 ):
     """Resolve pipeline variables.
@@ -251,8 +220,6 @@ def resolve_pipeline_variables(
         func_args: function args.
         func_kwargs: function kwargs.
         hmac_key: key used to encrypt serialized and deserialized function and arguments.
-        s3_base_uri: the s3 base uri of the function step that the serialized artifacts
-            will be uploaded to. The s3_base_uri = s3_root_uri + pipeline_name.
         **settings: settings to pass to the deserialization function.
     """
 
@@ -274,9 +241,7 @@ def resolve_pipeline_variables(
     delayed_return_resolver = _DelayedReturnResolver(
         delayed_returns=delayed_returns,
         hmac_key=hmac_key,
-        parameter_resolver=parameter_resolver,
-        execution_variable_resolver=execution_variable_resolver,
-        s3_base_uri=s3_base_uri,
+        properties_resolver=properties_resolver,
         **settings,
     )
 
@@ -322,39 +287,27 @@ 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
+    return converted_func_args, converted_func_kwargs
 
-    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,
-            )
+def _convert_pipeline_variable_to_pickleable(arg):
+    """Convert a pipeline variable to pickleable."""
+    from sagemaker.workflow.entities import PipelineVariable
 
-        if isinstance(arg, PipelineVariable):
-            return arg._pickleable
+    from sagemaker.workflow.function_step import DelayedReturn
 
-        return arg
+    if isinstance(arg, DelayedReturn):
+        return _DelayedReturn(
+            uri=get_step(arg)._properties.OutputDataConfig.S3OutputPath._pickleable,
+            reference_path=arg._reference_path,
+        )
 
-    converted_func_args = tuple(convert(arg) for arg in func_args)
-    converted_func_kwargs = {key: convert(arg) for key, arg in func_kwargs.items()}
+    if isinstance(arg, PipelineVariable):
+        return arg._pickleable
 
-    return converted_func_args, converted_func_kwargs
+    return arg

src/sagemaker/remote_function/core/stored_function.py

@@ -175,7 +175,6 @@ def load_and_invoke(self) -> Any:
             args,
             kwargs,
             hmac_key=self.hmac_key,
-            s3_base_uri=self.s3_base_uri,
             sagemaker_session=self.sagemaker_session,
         )
 

src/sagemaker/remote_function/job.py

@@ -59,7 +59,10 @@
 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,
+)
 from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
     RuntimeEnvironmentManager,
     _DependencySettings,
@@ -72,6 +75,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
@@ -701,6 +705,7 @@ def compile(
         """Build the artifacts and generate the training job request."""
         from sagemaker.workflow.properties import Properties
         from sagemaker.workflow.parameters import Parameter
+        from sagemaker.workflow.functions import Join
         from sagemaker.workflow.execution_variables import ExecutionVariables, ExecutionVariable
         from sagemaker.workflow.utilities import load_step_compilation_context
 
@@ -760,7 +765,19 @@ def compile(
             job_settings=job_settings, s3_base_uri=s3_base_uri
         )
 
-        output_config = {"S3OutputPath": s3_base_uri}
+        if step_compilation_context:
+            s3_output_path = Join(
+                on="/",
+                values=[
+                    s3_base_uri,
+                    ExecutionVariables.PIPELINE_EXECUTION_ID,
+                    step_compilation_context.step_name,
+                    "results",
+                ],
+            )
+            output_config = {"S3OutputPath": s3_output_path}
+        else:
+            output_config = {"S3OutputPath": s3_base_uri}
         if job_settings.s3_kms_key is not None:
             output_config["KmsKeyId"] = job_settings.s3_kms_key
         request_dict["OutputDataConfig"] = output_config
@@ -804,6 +821,11 @@ def compile(
                 if isinstance(arg, (Parameter, ExecutionVariable, Properties)):
                     container_args.extend([arg.expr["Get"], arg.to_string()])
 
+                if isinstance(arg, DelayedReturn):
+                    # The uri is a _Properties object
+                    uri = _convert_pipeline_variable_to_pickleable(arg).uri
+                    container_args.extend([uri.path, {"Get": uri.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
@@ -101,6 +102,12 @@ def __init__(
 
         self.__job_settings = None
 
+        # It's for internal usage to retrieve execution id from the properties.
+        # However, we won't expose the properties of function step to customers.
+        self._properties = Properties(
+            step_name=name, step=self, shape_name="DescribeTrainingJobResponse"
+        )
+
         (
             self._converted_func_args,
             self._converted_func_kwargs,

src/sagemaker/workflow/pipeline.py

@@ -28,7 +28,6 @@
 from sagemaker._studio import _append_project_tags
 from sagemaker.config import PIPELINE_ROLE_ARN_PATH, PIPELINE_TAGS_PATH
 from sagemaker.remote_function.core.serialization import deserialize_obj_from_s3
-from sagemaker.remote_function.core.stored_function import RESULTS_FOLDER
 from sagemaker.remote_function.errors import RemoteFunctionError
 from sagemaker.remote_function.job import JOBS_CONTAINER_ENTRYPOINT
 from sagemaker.s3_utils import s3_path_join
@@ -977,23 +976,20 @@ def result(self, step_name: str):
         return get_function_step_result(
             step_name=step_name,
             step_list=self.list_steps(),
-            execution_id=self.arn.split("/")[-1],
             sagemaker_session=self.sagemaker_session,
         )
 
 
 def get_function_step_result(
     step_name: str,
     step_list: list,
-    execution_id: str,
     sagemaker_session: Session,
 ):
     """Helper function to retrieve the output of a ``@step`` decorated function.
 
     Args:
         step_name (str): The name of the pipeline step.
         step_list (list): A list of executed pipeline steps of the specified execution.
-        execution_id (str): The specified id of the pipeline execution.
         sagemaker_session (Session): Session object which manages interactions
             with Amazon SageMaker APIs and any other AWS services needed.
     Returns:
@@ -1043,7 +1039,7 @@ def get_function_step_result(
     if job_status == "Completed":
         return deserialize_obj_from_s3(
             sagemaker_session=sagemaker_session,
-            s3_uri=s3_path_join(s3_output_path, execution_id, step_name, RESULTS_FOLDER),
+            s3_uri=s3_path_join(s3_output_path),
             hmac_key=describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"],
         )
 

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