fix: Move func and args serialization of function step to step level

Author: qidewenwhen

src/sagemaker/remote_function/core/pipeline_variables.py

@@ -30,6 +30,7 @@ class Context:
     property_references: Dict[str, str] = field(default_factory=dict)
     serialize_output_to_json: bool = False
     func_step_s3_dir: str = None
+    s3_base_uri: str = None
 
 
 @dataclass
@@ -77,6 +78,17 @@ 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."""
@@ -155,6 +167,7 @@ def __init__(
         hmac_key: str,
         parameter_resolver: _ParameterResolver,
         execution_variable_resolver: _ExecutionVariableResolver,
+        s3_base_uri: str,
         **settings,
     ):
         """Resolve delayed return.
@@ -164,8 +177,11 @@ def __init__(
             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 DelayedReturn object associates with.
             **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
         # different delayed returns can have the same uri, so we need to dedupe
@@ -205,6 +221,8 @@ def _resolve_delayed_return_uri(self, delayed_return: _DelayedReturn):
                 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)
@@ -251,6 +269,7 @@ def resolve_pipeline_variables(
         hmac_key=hmac_key,
         parameter_resolver=parameter_resolver,
         execution_variable_resolver=execution_variable_resolver,
+        s3_base_uri=context.s3_base_uri,
         **settings,
     )
 
@@ -289,11 +308,10 @@ def resolve_pipeline_variables(
     return resolved_func_args, resolved_func_kwargs
 
 
-def convert_pipeline_variables_to_pickleable(s3_base_uri: str, func_args: Tuple, func_kwargs: Dict):
+def convert_pipeline_variables_to_pickleable(func_args: Tuple, func_kwargs: Dict):
     """Convert pipeline variables to pickleable.
 
     Args:
-        s3_base_uri: s3 base uri where artifacts are stored.
         func_args: function args.
         func_kwargs: function kwargs.
     """
@@ -304,11 +322,19 @@ def convert_pipeline_variables_to_pickleable(s3_base_uri: str, func_args: Tuple,
 
     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 and after function. After step-level arg serialization,
+    # it's hard to update the s3_base_uri in pipeline compile time.
+    # Thus set a placeholder _S3BaseUriIdentifier here and let the runtime job to resolve it.
+    # 2. For saying s3_root_uri is unknown, it's because when defining function steps,
+    # the sagemaker_session is not passed in the pipeline 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=[
-                    s3_base_uri,
+                    _S3BaseUriIdentifier(),
                     ExecutionVariables.PIPELINE_EXECUTION_ID._pickleable,
                     arg._step.name,
                     "results",

src/sagemaker/remote_function/core/serialization.py

@@ -161,17 +161,13 @@ def serialize_func_to_s3(
     Raises:
         SerializationError: when fail to serialize function to bytes.
     """
-    bytes_to_upload = CloudpickleSerializer.serialize(func)
 
-    _upload_bytes_to_s3(bytes_to_upload, f"{s3_uri}/payload.pkl", s3_kms_key, sagemaker_session)
-
-    sha256_hash = _compute_hash(bytes_to_upload, secret_key=hmac_key)
-
-    _upload_bytes_to_s3(
-        _MetaData(sha256_hash).to_json(),
-        f"{s3_uri}/metadata.json",
-        s3_kms_key,
-        sagemaker_session,
+    _upload_payload_and_metadata_to_s3(
+        bytes_to_upload=CloudpickleSerializer.serialize(func),
+        hmac_key=hmac_key,
+        s3_uri=s3_uri,
+        sagemaker_session=sagemaker_session,
+        s3_kms_key=s3_kms_key,
     )
 
 
@@ -220,17 +216,12 @@ def serialize_obj_to_s3(
         SerializationError: when fail to serialize object to bytes.
     """
 
-    bytes_to_upload = CloudpickleSerializer.serialize(obj)
-
-    _upload_bytes_to_s3(bytes_to_upload, f"{s3_uri}/payload.pkl", s3_kms_key, sagemaker_session)
-
-    sha256_hash = _compute_hash(bytes_to_upload, secret_key=hmac_key)
-
-    _upload_bytes_to_s3(
-        _MetaData(sha256_hash).to_json(),
-        f"{s3_uri}/metadata.json",
-        s3_kms_key,
-        sagemaker_session,
+    _upload_payload_and_metadata_to_s3(
+        bytes_to_upload=CloudpickleSerializer.serialize(obj),
+        hmac_key=hmac_key,
+        s3_uri=s3_uri,
+        sagemaker_session=sagemaker_session,
+        s3_kms_key=s3_kms_key,
     )
 
 
@@ -318,8 +309,32 @@ def serialize_exception_to_s3(
     """
     pickling_support.install()
 
-    bytes_to_upload = CloudpickleSerializer.serialize(exc)
+    _upload_payload_and_metadata_to_s3(
+        bytes_to_upload=CloudpickleSerializer.serialize(exc),
+        hmac_key=hmac_key,
+        s3_uri=s3_uri,
+        sagemaker_session=sagemaker_session,
+        s3_kms_key=s3_kms_key,
+    )
 
+
+def _upload_payload_and_metadata_to_s3(
+    bytes_to_upload: Union[bytes, io.BytesIO],
+    hmac_key: str,
+    s3_uri: str,
+    sagemaker_session: Session,
+    s3_kms_key,
+):
+    """Uploads serialized payload and metadata to s3.
+
+    Args:
+        bytes_to_upload (bytes): Serialized bytes to upload.
+        hmac_key (str): Key used to calculate hmac-sha256 hash of the serialized obj.
+        s3_uri (str): S3 root uri to which resulting serialized artifacts will be uploaded.
+        sagemaker_session (sagemaker.session.Session):
+            The underlying Boto3 session which AWS service calls are delegated to.
+        s3_kms_key (str): KMS key used to encrypt artifacts uploaded to S3.
+    """
     _upload_bytes_to_s3(bytes_to_upload, f"{s3_uri}/payload.pkl", s3_kms_key, sagemaker_session)
 
     sha256_hash = _compute_hash(bytes_to_upload, secret_key=hmac_key)

src/sagemaker/remote_function/core/stored_function.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import os
+from dataclasses import dataclass
 from typing import Any
 
 
@@ -36,6 +37,14 @@
 JSON_RESULTS_FILE = "results.json"
 
 
+@dataclass
+class _SerializedData:
+    """Data class to store serialized function and arguments"""
+
+    func: bytes
+    args: bytes
+
+
 class StoredFunction:
     """Class representing a remote function stored in S3."""
 
@@ -105,6 +114,38 @@ def save(self, func, *args, **kwargs):
             s3_kms_key=self.s3_kms_key,
         )
 
+    def save_pipeline_step_function(self, serialized_data):
+        """Upload serialized function and arguments to s3.
+
+        Args:
+            serialized_data (_SerializedData): The serialized function
+                and function arguments of a function step.
+        """
+
+        logger.info(
+            "Uploading serialized function code to %s",
+            s3_path_join(self.func_upload_path, FUNCTION_FOLDER),
+        )
+        serialization._upload_payload_and_metadata_to_s3(
+            bytes_to_upload=serialized_data.func,
+            hmac_key=self.hmac_key,
+            s3_uri=s3_path_join(self.func_upload_path, FUNCTION_FOLDER),
+            sagemaker_session=self.sagemaker_session,
+            s3_kms_key=self.s3_kms_key,
+        )
+
+        logger.info(
+            "Uploading serialized function arguments to %s",
+            s3_path_join(self.func_upload_path, ARGUMENTS_FOLDER),
+        )
+        serialization._upload_payload_and_metadata_to_s3(
+            bytes_to_upload=serialized_data.args,
+            hmac_key=self.hmac_key,
+            s3_uri=s3_path_join(self.func_upload_path, ARGUMENTS_FOLDER),
+            sagemaker_session=self.sagemaker_session,
+            s3_kms_key=self.s3_kms_key,
+        )
+
     def load_and_invoke(self) -> Any:
         """Load and deserialize the function and the arguments and then execute it."""
 

src/sagemaker/remote_function/invoke_function.py

@@ -84,6 +84,7 @@ def _load_pipeline_context(args) -> Context:
     property_references = args.property_references
     serialize_output_to_json = args.serialize_output_to_json
     func_step_s3_dir = args.func_step_s3_dir
+    s3_base_uri = args.s3_base_uri
 
     property_references_dict = {}
     for i in range(0, len(property_references), 2):
@@ -94,6 +95,7 @@ def _load_pipeline_context(args) -> Context:
         property_references=property_references_dict,
         serialize_output_to_json=serialize_output_to_json,
         func_step_s3_dir=func_step_s3_dir,
+        s3_base_uri=s3_base_uri,
     )
 
 

src/sagemaker/remote_function/job.py

@@ -52,11 +52,8 @@
 from sagemaker.utils import name_from_base, _tmpdir, resolve_value_from_config
 from sagemaker.s3 import s3_path_join, S3Uploader
 from sagemaker import vpc_utils
-from sagemaker.remote_function.core.stored_function import StoredFunction
-from sagemaker.remote_function.core.pipeline_variables import (
-    Context,
-    convert_pipeline_variables_to_pickleable,
-)
+from sagemaker.remote_function.core.stored_function import StoredFunction, _SerializedData
+from sagemaker.remote_function.core.pipeline_variables import Context
 from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
     RuntimeEnvironmentManager,
     _DependencySettings,
@@ -695,6 +692,7 @@ def compile(
         func_args: tuple,
         func_kwargs: dict,
         run_info=None,
+        serialized_data: _SerializedData = None,
     ) -> dict:
         """Build the artifacts and generate the training job request."""
         from sagemaker.workflow.properties import Properties
@@ -732,12 +730,8 @@ def compile(
                     func_step_s3_dir=step_compilation_context.pipeline_build_time,
                 ),
             )
-            converted_func_args, converted_func_kwargs = convert_pipeline_variables_to_pickleable(
-                s3_base_uri=s3_base_uri,
-                func_args=func_args,
-                func_kwargs=func_kwargs,
-            )
-            stored_function.save(func, *converted_func_args, **converted_func_kwargs)
+
+            stored_function.save_pipeline_step_function(serialized_data)
 
         stopping_condition = {
             "MaxRuntimeInSeconds": job_settings.max_runtime_in_seconds,

src/sagemaker/workflow/function_step.py

@@ -83,6 +83,11 @@ def __init__(
             func_kwargs (dict): keyword arguments of the python function.
             **kwargs: Additional arguments to be passed to the `step` decorator.
         """
+        from sagemaker.remote_function.core.pipeline_variables import (
+            convert_pipeline_variables_to_pickleable,
+        )
+        from sagemaker.remote_function.core.serialization import CloudpickleSerializer
+        from sagemaker.remote_function.core.stored_function import _SerializedData
 
         super(_FunctionStep, self).__init__(
             name, StepTypeEnum.TRAINING, display_name, description, depends_on, retry_policies
@@ -96,6 +101,21 @@ def __init__(
 
         self.__job_settings = None
 
+        (
+            self._converted_func_args,
+            self._converted_func_kwargs,
+        ) = convert_pipeline_variables_to_pickleable(
+            func_args=self._func_args,
+            func_kwargs=self._func_kwargs,
+        )
+
+        self._serialized_data = _SerializedData(
+            func=CloudpickleSerializer.serialize(self._func),
+            args=CloudpickleSerializer.serialize(
+                (self._converted_func_args, self._converted_func_kwargs)
+            ),
+        )
+
     @property
     def func(self):
         """The python function to run as a pipeline step."""
@@ -185,6 +205,7 @@ def arguments(self) -> RequestType:
             func=self.func,
             func_args=self.func_args,
             func_kwargs=self.func_kwargs,
+            serialized_data=self._serialized_data,
         )
         # Continue to pop job name if not explicitly opted-in via config
         request_dict = trim_request_dict(request_dict, "TrainingJobName", step_compilation_context)

tests/integ/sagemaker/workflow/test_step_decorator.py

@@ -858,3 +858,61 @@ def cuberoot(x):
             pipeline.delete()
         except Exception:
             pass
+
+
+def test_step_level_serialization(
+    sagemaker_session, role, pipeline_name, region_name, dummy_container_without_error
+):
+    os.environ["AWS_DEFAULT_REGION"] = region_name
+
+    _EXPECTED_STEP_A_OUTPUT = "This pipeline is a function."
+    _EXPECTED_STEP_B_OUTPUT = "This generates a function arg."
+
+    step_config = dict(
+        role=role,
+        image_uri=dummy_container_without_error,
+        instance_type=INSTANCE_TYPE,
+    )
+
+    # This pipeline function may clash with the pipeline object
+    # defined below.
+    # However, if the function and args serialization happen in
+    # step level, this clash won't happen.
+    def pipeline():
+        return _EXPECTED_STEP_A_OUTPUT
+
+    @step(**step_config)
+    def generator():
+        return _EXPECTED_STEP_B_OUTPUT
+
+    @step(**step_config)
+    def func_with_collision(var: str):
+        return f"{pipeline()} {var}"
+
+    step_output_a = generator()
+    step_output_b = func_with_collision(step_output_a)
+
+    pipeline = Pipeline(  # noqa: F811
+        name=pipeline_name,
+        steps=[step_output_b],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        create_and_execute_pipeline(
+            pipeline=pipeline,
+            pipeline_name=pipeline_name,
+            region_name=region_name,
+            role=role,
+            no_of_steps=2,
+            last_step_name=get_step(step_output_b).name,
+            execution_parameters=dict(),
+            step_status="Succeeded",
+            step_result_type=str,
+            step_result_value=f"{_EXPECTED_STEP_A_OUTPUT} {_EXPECTED_STEP_B_OUTPUT}",
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass