Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions src/sagemaker/workflow/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,7 +238,10 @@ def _create_args(
}

update_args(
kwargs, PipelineDescription=description, ParallelismConfiguration=parallelism_config
kwargs,
PipelineDescription=description,
# boto expects a dict, not a ParallelismConfiguration object (see issue #5354).
ParallelismConfiguration=_resolve_parallelism_config(parallelism_config),
)
return kwargs

Expand Down Expand Up @@ -403,7 +406,8 @@ def start(
kwargs,
PipelineExecutionDescription=execution_description,
PipelineExecutionDisplayName=execution_display_name,
ParallelismConfiguration=parallelism_config,
# boto expects a dict, not a ParallelismConfiguration object (see issue #5354).
ParallelismConfiguration=_resolve_parallelism_config(parallelism_config),
SelectiveExecutionConfig=selective_execution_config,
PipelineVersionId=pipeline_version_id,
)
Expand Down Expand Up @@ -862,6 +866,20 @@ def _map_lambda_outputs(steps: List[Step]):
return lambda_output_map


def _resolve_parallelism_config(parallelism_config):
"""Normalize a parallelism_config into the request dict boto expects.

boto's create/update/start pipeline APIs expect ``ParallelismConfiguration`` as a dict
(``{"MaxParallelExecutionSteps": int}``), not a ``ParallelismConfiguration`` object
(issue #5354). This converts the object via ``to_request()``. A dict is passed through
unchanged so callers who adopted the pre-fix ``.to_request()`` workaround keep working,
and ``None`` is returned as-is so ``update_args`` can drop the key.
"""
if isinstance(parallelism_config, ParallelismConfiguration):
return parallelism_config.to_request()
return parallelism_config


def update_args(args: Dict[str, Any], **kwargs):
"""Updates the request arguments dict with a value, if populated.

Expand Down
96 changes: 96 additions & 0 deletions tests/unit/sagemaker/workflow/test_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,102 @@ def test_pipeline_create_and_start_with_parallelism_config(sagemaker_session_moc
)


def test_pipeline_create_converts_parallelism_config_object_to_dict(
sagemaker_session_mock, role_arn
):
"""Regression for #5354.

create() (via _create_args) must convert a ParallelismConfiguration object to its request
dict, because boto rejects the object with a ParamValidationError.
"""
from sagemaker.workflow.parallelism_config import ParallelismConfiguration

pipeline = Pipeline(
name="MyPipeline",
parameters=[],
steps=[],
sagemaker_session=sagemaker_session_mock,
)
pipeline.create(
role_arn=role_arn,
parallelism_config=ParallelismConfiguration(max_parallel_execution_steps=5),
)
call_kwargs = sagemaker_session_mock.sagemaker_client.create_pipeline.call_args[1]
assert call_kwargs["ParallelismConfiguration"] == {"MaxParallelExecutionSteps": 5}
assert isinstance(call_kwargs["ParallelismConfiguration"], dict)


def test_pipeline_start_converts_parallelism_config_object_to_dict(
sagemaker_session_mock, role_arn
):
"""Regression for #5354: start() must pass a dict to boto, not a ParallelismConfiguration."""
from sagemaker.workflow.parallelism_config import ParallelismConfiguration

sagemaker_session_mock.sagemaker_client.start_pipeline_execution.return_value = dict(
PipelineExecutionArn="pipeline-execution-arn"
)
pipeline = Pipeline(
name="MyPipeline",
parameters=[],
steps=[],
sagemaker_session=sagemaker_session_mock,
)
pipeline.start(parallelism_config=ParallelismConfiguration(max_parallel_execution_steps=7))
call_kwargs = sagemaker_session_mock.sagemaker_client.start_pipeline_execution.call_args[1]
assert call_kwargs["ParallelismConfiguration"] == {"MaxParallelExecutionSteps": 7}
assert isinstance(call_kwargs["ParallelismConfiguration"], dict)


def test_resolve_parallelism_config_output_passes_real_boto_param_validation():
"""End-to-end proof for #5354 against the real SageMaker service model (offline).

boto3's client-side ParamValidator (no network, no credentials) is exactly what raised the
reported ParamValidationError. This asserts the raw ParallelismConfiguration object is
rejected by that validator, while the value produced by ``_resolve_parallelism_config`` is
accepted -- reproducing the bug and proving the fix at the boto layer, not just via mocks.
"""
import boto3
from botocore import validate

from sagemaker.workflow.parallelism_config import ParallelismConfiguration
from sagemaker.workflow.pipeline import _resolve_parallelism_config

shape = (
boto3.client("sagemaker", region_name="us-west-2")
.meta.service_model.operation_model("CreatePipeline")
.input_shape
)
validator = validate.ParamValidator()
base = {
"PipelineName": "p",
"RoleArn": "arn:aws:iam::111111111111:role/x",
"PipelineDefinition": "{}",
"ClientRequestToken": "x" * 32,
}
cfg = ParallelismConfiguration(max_parallel_execution_steps=5)

# The raw object is what the old code passed -- boto rejects it.
raw_report = validator.validate({**base, "ParallelismConfiguration": cfg}, shape)
assert raw_report.has_errors()

# The converted value the fix passes -- boto accepts it.
fixed_report = validator.validate(
{**base, "ParallelismConfiguration": _resolve_parallelism_config(cfg)}, shape
)
assert not fixed_report.has_errors()


def test_resolve_parallelism_config_passes_through_dict_and_none():
"""Backwards compat for #5354: a dict (pre-fix workaround) is forwarded unchanged and
None is returned as-is so update_args drops the key."""
from sagemaker.workflow.pipeline import _resolve_parallelism_config

assert _resolve_parallelism_config({"MaxParallelExecutionSteps": 3}) == {
"MaxParallelExecutionSteps": 3
}
assert _resolve_parallelism_config(None) is None


@patch("sagemaker.s3.S3Uploader.upload_string_as_file_body")
def test_large_pipeline_create(sagemaker_session_mock, role_arn):
sagemaker_session_mock.sagemaker_config = {}
Expand Down
Loading