diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py index ee4bf562ce..05622e9ca0 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py @@ -281,7 +281,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 @@ -466,7 +469,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, MlflowExperimentName=mlflow_experiment_name, PipelineVersionId=pipeline_version_id, @@ -967,6 +971,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. diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py index 8ea04daab2..759b906c02 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py @@ -248,6 +248,99 @@ def test_pipeline_start_local_mode(mock_session, mock_step): mock_session.sagemaker_client.start_pipeline_execution.assert_called_once() +def test_create_args_converts_parallelism_config_to_dict(mock_session, mock_step): + """Regression for #5354. + + ``_create_args`` (used by create/update/upsert) must convert a ParallelismConfiguration + object to its request dict, because boto rejects the object with a ParamValidationError. + """ + from sagemaker.mlops.workflow.parallelism_config import ParallelismConfiguration + + pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session) + with patch.object(pipeline, "definition", return_value="{}"): + kwargs = pipeline._create_args( + "role-arn", "desc", ParallelismConfiguration(max_parallel_execution_steps=5) + ) + + assert kwargs["ParallelismConfiguration"] == {"MaxParallelExecutionSteps": 5} + assert isinstance(kwargs["ParallelismConfiguration"], dict) + + +def test_create_args_omits_parallelism_config_when_none(mock_session, mock_step): + """A None parallelism_config must not add the key (and must not raise).""" + pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session) + with patch.object(pipeline, "definition", return_value="{}"): + kwargs = pipeline._create_args("role-arn", "desc", None) + + assert "ParallelismConfiguration" not in kwargs + + +def test_create_args_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 that 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.mlops.workflow.parallelism_config import ParallelismConfiguration + from sagemaker.mlops.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_create_args_passes_through_dict_parallelism_config(mock_session, mock_step): + """Backwards compat for #5354: a caller who used the pre-fix ``.to_request()`` workaround + passes an already-converted dict; it must be forwarded unchanged, not re-converted.""" + pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session) + with patch.object(pipeline, "definition", return_value="{}"): + kwargs = pipeline._create_args("role-arn", "desc", {"MaxParallelExecutionSteps": 3}) + + assert kwargs["ParallelismConfiguration"] == {"MaxParallelExecutionSteps": 3} + + +def test_pipeline_start_converts_parallelism_config_to_dict(mock_session, mock_step): + """Regression for #5354: start() must pass a dict to boto, not a ParallelismConfiguration.""" + from sagemaker.mlops.workflow.parallelism_config import ParallelismConfiguration + + mock_session.local_mode = False + mock_session.sagemaker_client.start_pipeline_execution.return_value = { + "PipelineExecutionArn": "arn" + } + pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session) + + pipeline.start(parallelism_config=ParallelismConfiguration(max_parallel_execution_steps=7)) + + call_kwargs = mock_session.sagemaker_client.start_pipeline_execution.call_args.kwargs + assert call_kwargs["ParallelismConfiguration"] == {"MaxParallelExecutionSteps": 7} + assert isinstance(call_kwargs["ParallelismConfiguration"], dict) + + def test_pipeline_get_latest_execution_arn_none(mock_session, mock_step): pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session) mock_session.sagemaker_client.list_pipeline_executions.return_value = {