This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0b32613480 Increase ConflictException retries to 4 total (#36337)
0b32613480 is described below
commit 0b326134801c8f6e1e9ad685ffb20b899c85c9ec
Author: Niko Oliveira <[email protected]>
AuthorDate: Thu Dec 21 11:05:51 2023 -0800
Increase ConflictException retries to 4 total (#36337)
* Increase ConflictException retries to 4 total
We have seen a recent uptick in ConflictExceptions we receive from
SageMaker. It's not many, but enough to fail our system tests
unnecessarily.
Bumping the retires to 4 which I still think is very reasonable,
considering it's a very tight loop of 0.3 seconds sleep per retry.
* Update airflow/providers/amazon/aws/hooks/sagemaker.py
Co-authored-by: Andrey Anshin <[email protected]>
* Fix tests
---------
Co-authored-by: Andrey Anshin <[email protected]>
---
airflow/providers/amazon/aws/hooks/sagemaker.py | 2 +-
tests/providers/amazon/aws/hooks/test_sagemaker.py | 6 ++++--
2 files changed, 5 insertions(+), 3 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py
b/airflow/providers/amazon/aws/hooks/sagemaker.py
index c2e671d427..101d17f92b 100644
--- a/airflow/providers/amazon/aws/hooks/sagemaker.py
+++ b/airflow/providers/amazon/aws/hooks/sagemaker.py
@@ -1138,7 +1138,7 @@ class SageMakerHook(AwsBaseHook):
if check_interval is None:
check_interval = 10
- for retries in (2, 1, 0):
+ for retries in reversed(range(5)):
try:
self.conn.stop_pipeline_execution(PipelineExecutionArn=pipeline_exec_arn)
except ClientError as ce:
diff --git a/tests/providers/amazon/aws/hooks/test_sagemaker.py
b/tests/providers/amazon/aws/hooks/test_sagemaker.py
index a511eab8c8..39c9ed3c79 100644
--- a/tests/providers/amazon/aws/hooks/test_sagemaker.py
+++ b/tests/providers/amazon/aws/hooks/test_sagemaker.py
@@ -811,6 +811,8 @@ class TestSageMakerHook:
operation_name="empty",
)
mock_conn().stop_pipeline_execution.side_effect = [
+ conflict_error,
+ conflict_error,
conflict_error,
conflict_error,
None,
@@ -819,7 +821,7 @@ class TestSageMakerHook:
hook = SageMakerHook(aws_conn_id="aws_default")
hook.stop_pipeline(pipeline_exec_arn="test")
- assert mock_conn().stop_pipeline_execution.call_count == 3
+ assert mock_conn().stop_pipeline_execution.call_count == 5
@patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
def test_stop_pipeline_fails_if_all_retries_error(self, mock_conn):
@@ -833,7 +835,7 @@ class TestSageMakerHook:
with pytest.raises(ClientError) as raised_exception:
hook.stop_pipeline(pipeline_exec_arn="test")
- assert mock_conn().stop_pipeline_execution.call_count == 3
+ assert mock_conn().stop_pipeline_execution.call_count == 5
assert raised_exception.value == conflict_error
@patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
