diff --git a/airflow/providers/amazon/aws/operators/batch.py b/airflow/providers/amazon/aws/operators/batch.py index 30264d9f4e1a7..96da41f1870a4 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 @@ def __init__( 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 @@ def __init__( ) 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 @@ def submit_job(self, context: Context): "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 c6a923b51ec2d..6bb2e8fe7463b 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -95,6 +95,8 @@ def test_init(self): 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