feature: Added parallelism config to create/update/start pipeline methods

Author: Ameen Khan

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -77,6 +77,12 @@ Pipeline
 .. autoclass:: sagemaker.workflow.pipeline._PipelineExecution
     :members:
 
+Parallelism Configuration
+-------------------------
+
+.. autoclass:: sagemaker.workflow.parallelism_config.ParallelismConfiguration
+    :members:
+
 Pipeline Experiment Config
 --------------------------
 

src/sagemaker/workflow/parallelism_config.py

@@ -0,0 +1,34 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Pipeline Parallelism Configuration"""
+from __future__ import absolute_import
+from sagemaker.workflow.entities import RequestType
+
+
+class ParallelismConfiguration:
+    """Parallelism config for SageMaker pipeline."""
+
+    def __init__(self, max_parallel_execution_steps: int):
+        """Create a ParallelismConfiguration
+
+        Args:
+            max_parallel_execution_steps, int:
+                max number of steps which could be parallelized
+        """
+        self.max_parallel_execution_steps = max_parallel_execution_steps
+
+    def to_request(self) -> RequestType:
+        """Returns: the request structure."""
+        return {
+            "MaxParallelExecutionSteps": self.max_parallel_execution_steps,
+        }

src/sagemaker/workflow/pipeline.py

@@ -35,6 +35,7 @@
 from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.parameters import Parameter
 from sagemaker.workflow.pipeline_experiment_config import PipelineExperimentConfig
+from sagemaker.workflow.parallelism_config import ParallelismConfiguration
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.steps import Step
 from sagemaker.workflow.step_collections import StepCollection
@@ -95,6 +96,7 @@ def create(
         role_arn: str,
         description: str = None,
         tags: List[Dict[str, str]] = None,
+        parallelism_config: ParallelismConfiguration = None,
     ) -> Dict[str, Any]:
         """Creates a Pipeline in the Pipelines service.
 
@@ -103,25 +105,33 @@ def create(
             description (str): A description of the pipeline.
             tags (List[Dict[str, str]]): A list of {"Key": "string", "Value": "string"} dicts as
                 tags.
+            parallelism_config (Optional[Config for parallel steps, Parallelism configuration that
+                is applied to each of. the executions
 
         Returns:
             A response dict from the service.
         """
         tags = _append_project_tags(tags)
-
-        kwargs = self._create_args(role_arn, description)
+        kwargs = self._create_args(role_arn, description, parallelism_config)
         update_args(
             kwargs,
             Tags=tags,
         )
         return self.sagemaker_session.sagemaker_client.create_pipeline(**kwargs)
 
-    def _create_args(self, role_arn: str, description: str):
+    def _create_args(
+        self,
+        role_arn: str,
+        description: str,
+        parallelism_config: ParallelismConfiguration
+    ):
         """Constructs the keyword argument dict for a create_pipeline call.
 
         Args:
             role_arn (str): The role arn that is assumed by pipelines to create step artifacts.
             description (str): A description of the pipeline.
+            parallelism_config (Optional[ParallelismConfiguration]): Config for parallel steps, that
+                is applied to each of the executions.
 
         Returns:
             A keyword argument dict for calling create_pipeline.
@@ -134,7 +144,7 @@ def _create_args(self, role_arn: str, description: str):
 
         # If pipeline definition is large, upload to S3 bucket and
         # provide PipelineDefinitionS3Location to request instead.
-        if len(pipeline_definition.encode("utf-8")) < 1024*100:
+        if len(pipeline_definition.encode("utf-8")) < 1024 * 100:
             kwargs["PipelineDefinition"] = self.definition()
         else:
             desired_s3_uri = s3.s3_path_join(
@@ -153,6 +163,7 @@ def _create_args(self, role_arn: str, description: str):
         update_args(
             kwargs,
             PipelineDescription=description,
+            ParallelismConfiguration=parallelism_config
         )
         return kwargs
 
@@ -166,24 +177,32 @@ def describe(self) -> Dict[str, Any]:
         """
         return self.sagemaker_session.sagemaker_client.describe_pipeline(PipelineName=self.name)
 
-    def update(self, role_arn: str, description: str = None) -> Dict[str, Any]:
+    def update(
+        self,
+        role_arn: str,
+        description: str = None,
+        parallelism_config: ParallelismConfiguration = None,
+    ) -> Dict[str, Any]:
         """Updates a Pipeline in the Workflow service.
 
         Args:
             role_arn (str): The role arn that is assumed by pipelines to create step artifacts.
             description (str): A description of the pipeline.
+            parallelism_config (Optional[ParallelismConfiguration]): Config for parallel steps, that
+                is applied to each of the executions.
 
         Returns:
             A response dict from the service.
         """
-        kwargs = self._create_args(role_arn, description)
+        kwargs = self._create_args(role_arn, description, parallelism_config)
         return self.sagemaker_session.sagemaker_client.update_pipeline(**kwargs)
 
     def upsert(
         self,
         role_arn: str,
         description: str = None,
         tags: List[Dict[str, str]] = None,
+        parallelism_config: ParallelismConfiguration = None,
     ) -> Dict[str, Any]:
         """Creates a pipeline or updates it, if it already exists.
 
@@ -192,12 +211,14 @@ def upsert(
             description (str): A description of the pipeline.
             tags (List[Dict[str, str]]): A list of {"Key": "string", "Value": "string"} dicts as
                 tags.
+            parallelism_config (Optional[Config for parallel steps, Parallelism configuration that
+                is applied to each of. the executions
 
         Returns:
             response dict from service
         """
         try:
-            response = self.create(role_arn, description, tags)
+            response = self.create(role_arn, description, tags, parallelism_config)
         except ClientError as e:
             error = e.response["Error"]
             if (
@@ -235,6 +256,7 @@ def start(
         parameters: Dict[str, Union[str, bool, int, float]] = None,
         execution_display_name: str = None,
         execution_description: str = None,
+        parallelism_config: ParallelismConfiguration = None,
     ):
         """Starts a Pipeline execution in the Workflow service.
 
@@ -243,6 +265,8 @@ def start(
                 pipeline parameters.
             execution_display_name (str): The display name of the pipeline execution.
             execution_description (str): A description of the execution.
+            parallelism_config (Optional[ParallelismConfiguration]): Config for parallel steps, that
+                is applied to each of the executions.
 
         Returns:
             A `_PipelineExecution` instance, if successful.
@@ -265,6 +289,7 @@ def start(
             PipelineParameters=format_start_parameters(parameters),
             PipelineExecutionDescription=execution_description,
             PipelineExecutionDisplayName=execution_display_name,
+            ParallelismConfiguration=parallelism_config,
         )
         response = self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs)
         return _PipelineExecution(

tests/integ/test_workflow.py

@@ -71,6 +71,7 @@
 )
 from sagemaker.workflow.step_collections import RegisterModel
 from sagemaker.workflow.pipeline import Pipeline
+from sagemaker.workflow.parallelism_config import ParallelismConfiguration
 from sagemaker.lambda_helper import Lambda
 from sagemaker.feature_store.feature_group import FeatureGroup, FeatureDefinition, FeatureTypeEnum
 from tests.integ import DATA_DIR
@@ -2290,7 +2291,7 @@ def test_large_pipeline(sagemaker_session, role, pipeline_name, region_name):
             sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
             inputs={"arg1": "foo"},
             outputs=[outputParam],
-        ) for count in range(500)
+        ) for count in range(2000)
     ]
     pipeline = Pipeline(
         name=pipeline_name,
@@ -2319,3 +2320,45 @@ def test_large_pipeline(sagemaker_session, role, pipeline_name, region_name):
             pipeline.delete()
         except Exception:
             pass
+
+def test_create_parallelism_config(sagemaker_session, role, pipeline_name, region_name):
+    instance_count = ParameterInteger(name="InstanceCount", default_value=2)
+
+    outputParam = CallbackOutput(output_name="output", output_type=CallbackOutputTypeEnum.String)
+
+    callback_steps = [
+        CallbackStep(
+            name=f"callback-step{count}",
+            sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+            inputs={"arg1": "foo"},
+            outputs=[outputParam],
+        )
+        for count in range(500)
+    ]
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_count],
+        steps=callback_steps,
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        response = pipeline.create(role, parallelism_config={"MaxParallelExecutionSteps": 50})
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+        pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
+        response = pipeline.update(role, parallelism_config={"MaxParallelExecutionSteps": 50})
+        update_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            update_arn,
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass