feature: method to build pipeline parameters from existing execution … (aws#951)

* feature: method to build pipeline parameters from existing execution with optional value overrides

* fix style check

* assert error message in unit test

Author: shenlongtang

src/sagemaker/workflow/pipeline.py

@@ -463,6 +463,76 @@ def _get_latest_execution_arn(self):
             return response["PipelineExecutionSummaries"][0]["PipelineExecutionArn"]
         return None
 
+    def build_parameters_from_execution(
+        self,
+        pipeline_execution_arn: str,
+        parameter_value_overrides: Dict[str, Union[str, bool, int, float]] = None,
+    ) -> Dict[str, Union[str, bool, int, float]]:
+        """Gets the parameters from an execution, update with optional parameter value overrides.
+
+        Args:
+            pipeline_execution_arn (str): The arn of the pipeline execution.
+            parameter_value_overrides (Dict[str, Union[str, bool, int, float]]): Parameter dict
+                to be updated in the parameters from the referenced execution.
+
+        Returns:
+            A parameter dict built from an execution and provided parameter value overrides.
+        """
+        execution_parameters = self._get_parameters_for_execution(pipeline_execution_arn)
+        if parameter_value_overrides is not None:
+            self._validate_parameter_overrides(
+                pipeline_execution_arn, execution_parameters, parameter_value_overrides
+            )
+            execution_parameters.update(parameter_value_overrides)
+        return execution_parameters
+
+    def _get_parameters_for_execution(self, pipeline_execution_arn: str) -> Dict[str, str]:
+        """Gets all the parameters from an execution.
+
+        Args:
+            pipeline_execution_arn (str): The arn of the pipeline execution.
+
+        Returns:
+            A parameter dict from the execution.
+        """
+        pipeline_execution = _PipelineExecution(
+            arn=pipeline_execution_arn,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        response = pipeline_execution.list_parameters()
+        parameter_list = response["PipelineParameters"]
+        while response.get("NextToken") is not None:
+            response = pipeline_execution.list_parameters(next_token=response["NextToken"])
+            parameter_list.extend(response["PipelineParameters"])
+
+        return {parameter["Name"]: parameter["Value"] for parameter in parameter_list}
+
+    @staticmethod
+    def _validate_parameter_overrides(
+        pipeline_execution_arn: str,
+        execution_parameters: Dict[str, str],
+        parameter_overrides: Dict[str, Union[str, bool, int, float]],
+    ):
+        """Validates the parameter overrides are present in the execution parameters.
+
+        Args:
+            pipeline_execution_arn (str): The arn of the pipeline execution.
+            execution_parameters (Dict[str, str]): A parameter dict from the execution.
+            parameter_overrides (Dict[str, Union[str, bool, int, float]]): Parameter dict to be
+                updated in the parameters from the referenced execution.
+
+        Raises:
+            ValueError: If any parameters in parameter overrides is not present in the
+                execution parameters.
+        """
+        invalid_parameters = set(parameter_overrides) - set(execution_parameters)
+        if invalid_parameters:
+            raise ValueError(
+                f"The following parameter overrides provided: {str(invalid_parameters)} "
+                + f"are not present in the pipeline execution: {pipeline_execution_arn}"
+            )
+
 
 def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]:
     """Formats start parameter overrides as a list of dicts.
@@ -652,6 +722,30 @@ def list_steps(self):
         )
         return response["PipelineExecutionSteps"]
 
+    def list_parameters(self, max_results: int = None, next_token: str = None):
+        """Gets a list of parameters for a pipeline execution.
+
+        Args:
+            max_results (int): The maximum number of parameters to return in the response.
+            next_token (str):  If the result of the previous ListPipelineParametersForExecution
+                request was truncated, the response includes a NextToken. To retrieve the next
+                set of parameters, use the token in the next request.
+
+        Returns:
+            Information about the parameters of the pipeline execution.
+            See boto3 client list_pipeline_executions
+            https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_pipeline_parameters_for_execution
+        """
+        kwargs = dict(PipelineExecutionArn=self.arn)
+        update_args(
+            kwargs,
+            MaxResults=max_results,
+            NextToken=next_token,
+        )
+        return self.sagemaker_session.sagemaker_client.list_pipeline_parameters_for_execution(
+            **kwargs
+        )
+
     def wait(self, delay=30, max_attempts=60):
         """Waits for a pipeline execution.
 

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -17,7 +17,7 @@
 
 import pytest
 
-from mock import Mock, patch
+from mock import Mock, call, patch
 
 from sagemaker import s3
 from sagemaker.session_settings import SessionSettings
@@ -718,13 +718,99 @@ def test_pipeline_list_executions(sagemaker_session_mock):
     assert executions["NextToken"] == "token"
 
 
+def test_pipeline_build_parameters_from_execution(sagemaker_session_mock):
+    pipeline = Pipeline(
+        name="MyPipeline",
+        sagemaker_session=sagemaker_session_mock,
+    )
+    reference_execution_arn = "reference_execution_arn"
+    parameter_value_overrides = {"TestParameterName": "NewParameterValue"}
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.return_value = {
+        "PipelineParameters": [{"Name": "TestParameterName", "Value": "TestParameterValue"}]
+    }
+    parameters = pipeline.build_parameters_from_execution(
+        pipeline_execution_arn=reference_execution_arn,
+        parameter_value_overrides=parameter_value_overrides,
+    )
+    assert (
+        sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.called_with(
+            PipelineExecutionArn=reference_execution_arn
+        )
+    )
+    assert len(parameters) == 1
+    assert parameters["TestParameterName"] == "NewParameterValue"
+
+
+def test_pipeline_build_parameters_from_execution_with_invalid_overrides(sagemaker_session_mock):
+    pipeline = Pipeline(
+        name="MyPipeline",
+        sagemaker_session=sagemaker_session_mock,
+    )
+    reference_execution_arn = "reference_execution_arn"
+    invalid_parameter_value_overrides = {"InvalidParameterName": "Value"}
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.return_value = {
+        "PipelineParameters": [{"Name": "TestParameterName", "Value": "TestParameterValue"}]
+    }
+    with pytest.raises(ValueError) as error:
+        pipeline.build_parameters_from_execution(
+            pipeline_execution_arn=reference_execution_arn,
+            parameter_value_overrides=invalid_parameter_value_overrides,
+        )
+    assert (
+        f"The following parameter overrides provided: {str(set(invalid_parameter_value_overrides.keys()))} "
+        + f"are not present in the pipeline execution: {reference_execution_arn}"
+        in str(error)
+    )
+    assert (
+        sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.called_with(
+            PipelineExecutionArn=reference_execution_arn
+        )
+    )
+
+
+def test_pipeline_build_parameters_from_execution_with_paginated_result(sagemaker_session_mock):
+    pipeline = Pipeline(
+        name="MyPipeline",
+        sagemaker_session=sagemaker_session_mock,
+    )
+    reference_execution_arn = "reference_execution_arn"
+    next_token = "token"
+    first_page_response = {
+        "PipelineParameters": [{"Name": "TestParameterName1", "Value": "TestParameterValue1"}],
+        "NextToken": next_token,
+    }
+    second_page_response = {
+        "PipelineParameters": [{"Name": "TestParameterName2", "Value": "TestParameterValue2"}],
+    }
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.side_effect = [
+        first_page_response,
+        second_page_response,
+    ]
+    parameters = pipeline.build_parameters_from_execution(
+        pipeline_execution_arn=reference_execution_arn
+    )
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.assert_has_calls(
+        [
+            call(PipelineExecutionArn=reference_execution_arn),
+            call(PipelineExecutionArn=reference_execution_arn, NextToken=next_token),
+        ]
+    )
+    assert len(parameters) == 2
+    assert parameters["TestParameterName1"] == "TestParameterValue1"
+    assert parameters["TestParameterName2"] == "TestParameterValue2"
+
+
 def test_pipeline_execution_basics(sagemaker_session_mock):
     sagemaker_session_mock.sagemaker_client.start_pipeline_execution.return_value = {
         "PipelineExecutionArn": "my:arn"
     }
     sagemaker_session_mock.sagemaker_client.list_pipeline_execution_steps.return_value = {
         "PipelineExecutionSteps": [Mock()]
     }
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.return_value = {
+        "PipelineParameters": [{"Name": "TestParameterName", "Value": "TestParameterValue"}],
+        "NextToken": "token",
+    }
     pipeline = Pipeline(
         name="MyPipeline",
         parameters=[ParameterString("alpha", "beta"), ParameterString("gamma", "delta")],
@@ -745,6 +831,17 @@ def test_pipeline_execution_basics(sagemaker_session_mock):
         PipelineExecutionArn="my:arn"
     )
     assert len(steps) == 1
+    list_parameters_response = execution.list_parameters()
+    assert (
+        sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.called_with(
+            PipelineExecutionArn="my:arn"
+        )
+    )
+    parameter_list = list_parameters_response["PipelineParameters"]
+    assert len(parameter_list) == 1
+    assert parameter_list[0]["Name"] == "TestParameterName"
+    assert parameter_list[0]["Value"] == "TestParameterValue"
+    assert list_parameters_response["NextToken"] == "token"
 
 
 def _generate_large_pipeline_steps(input_data: object):