Description

Fix build failures due to Sagemaker ThrottlingException when running pipeline integration tests

Fixes #(issue) -N/A

Why is the change necessary?

Recent build failures were due to Sagemaker ThrottlingException (Rate exceeded) during following tests:

• test_training_pipeline_estimators.py::test_pca_estimator
• test_training_pipeline_framework_estimator.py::test_torch_training_pipeline
• test_inference_pipeline.py::test_inference_pipeline_framework

Solution

Add an optional retry argument to the pipeline template constructors (InferencePipeline andTrainingPipeline) in order to add a retry strategy for each pipeline steps. The same retrier will be added for each step.

Caveat: This fix applies the retry strategy to all steps in the pipeline. The customer won't be able to customize the strategy for each step.

Alternate solution 1:

We could add the option for the client to customize retry strategies for each pipeline step by accepting adict, in addition to acceptingRetry object.

Caveat: The retry strategy dict keys must correspond exactly to the step variable names - A validation step could be added to warn the customer of any unrecognized keys.

For example:

retry_strategy_per_step = {   'training_step': <training_retry_strategy>,   'model_step': <model_retry_strategy>,   'endpoint_config_step': <endpoint_config_retry_strategy>,   'deploy_step': <deploy_retry_strategy>}

If adict is received, only add retriers to steps with defined strategies in that dict.

Alternate solution 2:

Only add retries to integration tests by updating the pipeline workflow with the added retries

# Once pipeline is created do something like:sagemaker_retry_strategy = Retry(    error_equals=["SageMaker.AmazonSageMakerException"],    interval_seconds=5,    max_attempts=5,    backoff_rate=2)steps = pipeline.workflow.definition.branch.stepsfor step in steps:    step.add_retry(sagemaker_retry_strategy)pipeline.workflow.update(Chain(steps))

Caveat: If the fix is only applied to the integration tests, customers who want to add retry strategies to the pipeline steps will have to do this each time they are creating a pipeline

Testing

• Updated integ test and added unit test
• Generated doc locally

Pull Request Checklist

Please check all boxes (including N/A items)

Testing

• Unit tests added
• Integration test added
• Manual testing - why was it necessary? could it be automated? -N/A

Documentation

• docs: All relevantdocs updated
• docstrings: All public APIs documented

Title and description

• Change type: Title is prefixed with change type: and followsconventional commits
• References: Indicate issues fixed via:Fixes #xxx -N/A

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license.