feature: Updated create/update pipeline to support large pipeline defintion

Author: Ameen Khan

src/sagemaker/workflow/pipeline.py

@@ -22,6 +22,7 @@
 import botocore
 from botocore.exceptions import ClientError
 
+from sagemaker import s3
 from sagemaker._studio import _append_project_tags
 from sagemaker.session import Session
 from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep
@@ -125,11 +126,30 @@ def _create_args(self, role_arn: str, description: str):
         Returns:
             A keyword argument dict for calling create_pipeline.
         """
+        pipeline_definition = self.definition()
         kwargs = dict(
             PipelineName=self.name,
-            PipelineDefinition=self.definition(),
             RoleArn=role_arn,
         )
+
+        # If pipeline definition is large, upload to S3 bucket and
+        # provide PipelineDefinitionS3Location to request instead.
+        if len(pipeline_definition.encode("utf-8")) < 1024*100:
+            kwargs["PipelineDefinition"] = self.definition()
+        else:
+            desired_s3_uri = s3.s3_path_join(
+                "s3://", self.sagemaker_session.default_bucket(), self.name
+            )
+            s3.S3Uploader.upload_string_as_file_body(
+                body=pipeline_definition,
+                desired_s3_uri=desired_s3_uri,
+                sagemaker_session=self.sagemaker_session,
+            )
+            kwargs["PipelineDefinitionS3Location"] = {
+                "Bucket": self.sagemaker_session.default_bucket(),
+                "ObjectKey": self.name,
+            }
+
         update_args(
             kwargs,
             PipelineDescription=description,

tests/integ/test_workflow.py

@@ -2757,3 +2757,45 @@ def cleanup_feature_group(feature_group: FeatureGroup):
         except Exception as e:
             print(f"Delete FeatureGroup failed with error: {e}.")
             pass
+
+
+def test_large_pipeline(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)
+        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)
+        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