This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-8-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 0e6b5d8f8e71df326d237c2028940d5cbf3bac2c
Author: Evgeny <evgenys...@users.noreply.github.com>
AuthorDate: Wed Nov 22 04:26:42 2023 -0500
Added retry strategy parameter to Amazon AWS provider Batch Operator to
allow dynamic Batch retry strategies (#35789)
---
airflow/providers/amazon/aws/operators/batch.py | 5 +++++
tests/providers/amazon/aws/operators/test_batch.py | 6 ++++++
2 files changed, 11 insertions(+)
diff --git a/airflow/providers/amazon/aws/operators/batch.py
b/airflow/providers/amazon/aws/operators/batch.py
index 7cf2cbfeac..e917d1d81d 100644
--- a/airflow/providers/amazon/aws/operators/batch.py
+++ b/airflow/providers/amazon/aws/operators/batch.py
@@ -112,6 +112,7 @@ class BatchOperator(BaseOperator):
"array_properties",
"node_overrides",
"parameters",
+ "retry_strategy",
"waiters",
"tags",
"wait_for_completion",
@@ -122,6 +123,7 @@ class BatchOperator(BaseOperator):
"container_overrides": "json",
"parameters": "json",
"node_overrides": "json",
+ "retry_strategy": "json",
}
@property
@@ -160,6 +162,7 @@ class BatchOperator(BaseOperator):
share_identifier: str | None = None,
scheduling_priority_override: int | None = None,
parameters: dict | None = None,
+ retry_strategy: dict | None = None,
job_id: str | None = None,
waiters: Any | None = None,
max_retries: int = 4200,
@@ -201,6 +204,7 @@ class BatchOperator(BaseOperator):
self.scheduling_priority_override = scheduling_priority_override
self.array_properties = array_properties
self.parameters = parameters or {}
+ self.retry_strategy = retry_strategy or {}
self.waiters = waiters
self.tags = tags or {}
self.wait_for_completion = wait_for_completion
@@ -287,6 +291,7 @@ class BatchOperator(BaseOperator):
"tags": self.tags,
"containerOverrides": self.container_overrides,
"nodeOverrides": self.node_overrides,
+ "retryStrategy": self.retry_strategy,
"shareIdentifier": self.share_identifier,
"schedulingPriorityOverride": self.scheduling_priority_override,
}
diff --git a/tests/providers/amazon/aws/operators/test_batch.py
b/tests/providers/amazon/aws/operators/test_batch.py
index 8eb6601dfd..8a0d0e788a 100644
--- a/tests/providers/amazon/aws/operators/test_batch.py
+++ b/tests/providers/amazon/aws/operators/test_batch.py
@@ -63,6 +63,7 @@ class TestBatchOperator:
max_retries=self.MAX_RETRIES,
status_retries=self.STATUS_RETRIES,
parameters=None,
+ retry_strategy=None,
container_overrides={},
array_properties=None,
aws_conn_id="airflow_test",
@@ -96,6 +97,7 @@ class TestBatchOperator:
assert self.batch.hook.max_retries == self.MAX_RETRIES
assert self.batch.hook.status_retries == self.STATUS_RETRIES
assert self.batch.parameters == {}
+ assert self.batch.retry_strategy == {}
assert self.batch.container_overrides == {}
assert self.batch.array_properties is None
assert self.batch.node_overrides is None
@@ -119,6 +121,7 @@ class TestBatchOperator:
"array_properties",
"node_overrides",
"parameters",
+ "retry_strategy",
"waiters",
"tags",
"wait_for_completion",
@@ -143,6 +146,7 @@ class TestBatchOperator:
containerOverrides={},
jobDefinition="hello-world",
parameters={},
+ retryStrategy={},
tags={},
)
@@ -166,6 +170,7 @@ class TestBatchOperator:
containerOverrides={},
jobDefinition="hello-world",
parameters={},
+ retryStrategy={},
tags={},
)
@@ -232,6 +237,7 @@ class TestBatchOperator:
"jobName": JOB_NAME,
"jobDefinition": "hello-world",
"parameters": {},
+ "retryStrategy": {},
"tags": {},
}
if override == "overrides":
