Add missing unit tests for s3 upload

Author: Zhankuil

src/sagemaker/workflow/pipeline.py

@@ -144,7 +144,7 @@ def _create_args(
         # 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()
+            kwargs["PipelineDefinition"] = pipeline_definition
         else:
             desired_s3_uri = s3.s3_path_join(
                 "s3://", self.sagemaker_session.default_bucket(), self.name

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -21,9 +21,11 @@
 
 from mock import Mock
 
+from sagemaker import s3
 from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline import Pipeline
+from sagemaker.workflow.parallelism_config import ParallelismConfiguration
 from sagemaker.workflow.pipeline_experiment_config import (
     PipelineExperimentConfig,
     PipelineExperimentConfigProperties,
@@ -62,7 +64,9 @@ def role_arn():
 
 @pytest.fixture
 def sagemaker_session_mock():
-    return Mock()
+    session_mock = Mock()
+    session_mock.default_bucket = Mock(name="default_bucket", return_value="s3_bucket")
+    return session_mock
 
 
 def test_pipeline_create(sagemaker_session_mock, role_arn):
@@ -78,6 +82,50 @@ def test_pipeline_create(sagemaker_session_mock, role_arn):
     )
 
 
+def test_pipeline_create_with_parallelism_config(sagemaker_session_mock, role_arn):
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[],
+        steps=[],
+        pipeline_experiment_config=ParallelismConfiguration(max_parallel_execution_steps=10),
+        sagemaker_session=sagemaker_session_mock,
+    )
+    pipeline.create(role_arn=role_arn)
+    assert sagemaker_session_mock.sagemaker_client.create_pipeline.called_with(
+        PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn,
+        ParallelismConfiguration={
+            "MaxParallelExecutionSteps": 10
+        }
+    )
+
+
+def test_large_pipeline_create(sagemaker_session_mock, role_arn):
+    parameter = ParameterString("MyStr")
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[parameter],
+        steps=[CustomStep(name="MyStep", input_data=parameter)] * 2000,
+        sagemaker_session=sagemaker_session_mock,
+    )
+
+    s3.S3Uploader.upload_string_as_file_body = Mock()
+
+    pipeline.create(role_arn=role_arn)
+
+    assert s3.S3Uploader.upload_string_as_file_body.called_with(
+        body=pipeline.definition(),
+        s3_uri="s3://s3_bucket/MyPipeline")
+
+    assert sagemaker_session_mock.sagemaker_client.create_pipeline.called_with(
+        PipelineName="MyPipeline",
+        PipelineDefinitionS3Location={
+            "Bucket": "s3_bucket",
+            "ObjectKey": "MyPipeline"
+        },
+        RoleArn=role_arn
+    )
+
+
 def test_pipeline_update(sagemaker_session_mock, role_arn):
     pipeline = Pipeline(
         name="MyPipeline",
@@ -91,6 +139,50 @@ def test_pipeline_update(sagemaker_session_mock, role_arn):
     )
 
 
+def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_arn):
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[],
+        steps=[],
+        pipeline_experiment_config=ParallelismConfiguration(max_parallel_execution_steps=10),
+        sagemaker_session=sagemaker_session_mock,
+    )
+    pipeline.create(role_arn=role_arn)
+    assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
+        PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn,
+        ParallelismConfiguration={
+            "MaxParallelExecutionSteps": 10
+        }
+    )
+
+
+def test_large_pipeline_update(sagemaker_session_mock, role_arn):
+    parameter = ParameterString("MyStr")
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[parameter],
+        steps=[CustomStep(name="MyStep", input_data=parameter)] * 2000,
+        sagemaker_session=sagemaker_session_mock,
+    )
+
+    s3.S3Uploader.upload_string_as_file_body = Mock()
+
+    pipeline.create(role_arn=role_arn)
+
+    assert s3.S3Uploader.upload_string_as_file_body.called_with(
+        body=pipeline.definition(),
+        s3_uri="s3://s3_bucket/MyPipeline")
+
+    assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
+        PipelineName="MyPipeline",
+        PipelineDefinitionS3Location={
+            "Bucket": "s3_bucket",
+            "ObjectKey": "MyPipeline"
+        },
+        RoleArn=role_arn
+    )
+
+
 def test_pipeline_upsert(sagemaker_session_mock, role_arn):
     sagemaker_session_mock.side_effect = [
         ClientError(