This is an automated email from the ASF dual-hosted git repository.
onikolas 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 612676b975 Support `shareIdentifier` in BatchOperator (#30829)
612676b975 is described below
commit 612676b975a2ff26541bb2581fbdf2befc6c3de9
Author: Phani Kumar <[email protected]>
AuthorDate: Sat Apr 29 03:34:16 2023 +0530
Support `shareIdentifier` in BatchOperator (#30829)
* Add basic test for shareIdentifier field
---
airflow/providers/amazon/aws/operators/batch.py | 11 +++++++++++
tests/providers/amazon/aws/operators/test_batch.py | 2 ++
2 files changed, 13 insertions(+)
diff --git a/airflow/providers/amazon/aws/operators/batch.py
b/airflow/providers/amazon/aws/operators/batch.py
index 30264d9f4e..96da41f187 100644
--- a/airflow/providers/amazon/aws/operators/batch.py
+++ b/airflow/providers/amazon/aws/operators/batch.py
@@ -58,6 +58,11 @@ class BatchOperator(BaseOperator):
:param overrides: DEPRECATED, use container_overrides instead with the
same value.
:param container_overrides: the `containerOverrides` parameter for boto3
(templated)
:param node_overrides: the `nodeOverrides` parameter for boto3 (templated)
+ :param share_identifier: The share identifier for the job. Don't specify
this parameter if the job queue
+ doesn't have a scheduling policy.
+ :param scheduling_priority_override: The scheduling priority for the job.
+ Jobs with a higher scheduling priority are scheduled before jobs with
a lower scheduling priority.
+ This overrides any scheduling priority in the job definition
:param array_properties: the `arrayProperties` parameter for boto3
:param parameters: the `parameters` for boto3 (templated)
:param job_id: the job ID, usually unknown (None) until the
@@ -126,6 +131,8 @@ class BatchOperator(BaseOperator):
container_overrides: dict | None = None,
array_properties: dict | None = None,
node_overrides: dict | None = None,
+ share_identifier: str | None = None,
+ scheduling_priority_override: int | None = None,
parameters: dict | None = None,
job_id: str | None = None,
waiters: Any | None = None,
@@ -161,6 +168,8 @@ class BatchOperator(BaseOperator):
)
self.node_overrides = node_overrides
+ self.share_identifier = share_identifier
+ self.scheduling_priority_override = scheduling_priority_override
self.array_properties = array_properties
self.parameters = parameters or {}
self.waiters = waiters
@@ -227,6 +236,8 @@ class BatchOperator(BaseOperator):
"tags": self.tags,
"containerOverrides": self.container_overrides,
"nodeOverrides": self.node_overrides,
+ "shareIdentifier": self.share_identifier,
+ "schedulingPriorityOverride": self.scheduling_priority_override,
}
try:
diff --git a/tests/providers/amazon/aws/operators/test_batch.py
b/tests/providers/amazon/aws/operators/test_batch.py
index c6a923b51e..6bb2e8fe74 100644
--- a/tests/providers/amazon/aws/operators/test_batch.py
+++ b/tests/providers/amazon/aws/operators/test_batch.py
@@ -95,6 +95,8 @@ class TestBatchOperator:
assert self.batch.container_overrides == {}
assert self.batch.array_properties is None
assert self.batch.node_overrides is None
+ assert self.batch.share_identifier is None
+ assert self.batch.scheduling_priority_override is None
assert self.batch.hook.region_name == "eu-west-1"
assert self.batch.hook.aws_conn_id == "airflow_test"
assert self.batch.hook.client == self.client_mock
