ECS Overrides for AWS Batch submit_job (#39903)

airflow/providers/amazon/aws/hooks/batch_client.py

@@ -102,6 +102,7 @@ def submit_job(
         arrayProperties: dict,
         parameters: dict,
         containerOverrides: dict,
+        ecsPropertiesOverride: dict,
         tags: dict,
     ) -> dict:
         """
@@ -119,6 +120,8 @@ def submit_job(
 
         :param containerOverrides: the same parameter that boto3 will receive
 
+        :param ecsPropertiesOverride: the same parameter that boto3 will receive
+
         :param tags: the same parameter that boto3 will receive
 
         :return: an API response

airflow/providers/amazon/aws/operators/batch.py

@@ -65,6 +65,7 @@ class BatchOperator(BaseOperator):
     :param job_queue: the queue name on AWS Batch
     :param overrides: DEPRECATED, use container_overrides instead with the same value.
     :param container_overrides: the `containerOverrides` parameter for boto3 (templated)
+    :param ecs_properties_override: the `ecsPropertiesOverride` 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.
@@ -112,6 +113,7 @@ class BatchOperator(BaseOperator):
         "job_queue",
         "container_overrides",
         "array_properties",
+        "ecs_properties_override",
         "node_overrides",
         "parameters",
         "retry_strategy",
@@ -124,6 +126,7 @@ class BatchOperator(BaseOperator):
     template_fields_renderers = {
         "container_overrides": "json",
         "parameters": "json",
+        "ecs_properties_override": "json",
         "node_overrides": "json",
         "retry_strategy": "json",
     }
@@ -160,6 +163,7 @@ def __init__(
         overrides: dict | None = None,  # deprecated
         container_overrides: dict | None = None,
         array_properties: dict | None = None,
+        ecs_properties_override: dict | None = None,
         node_overrides: dict | None = None,
         share_identifier: str | None = None,
         scheduling_priority_override: int | None = None,
@@ -201,6 +205,7 @@ def __init__(
                 stacklevel=2,
             )
 
+        self.ecs_properties_override = ecs_properties_override
         self.node_overrides = node_overrides
         self.share_identifier = share_identifier
         self.scheduling_priority_override = scheduling_priority_override
@@ -296,6 +301,8 @@ def submit_job(self, context: Context):
             self.log.info("AWS Batch job - container overrides: %s", self.container_overrides)
         if self.array_properties:
             self.log.info("AWS Batch job - array properties: %s", self.array_properties)
+        if self.ecs_properties_override:
+            self.log.info("AWS Batch job - ECS properties: %s", self.ecs_properties_override)
         if self.node_overrides:
             self.log.info("AWS Batch job - node properties: %s", self.node_overrides)
 
@@ -307,6 +314,7 @@ def submit_job(self, context: Context):
             "parameters": self.parameters,
             "tags": self.tags,
             "containerOverrides": self.container_overrides,
+            "ecsPropertiesOverride": self.ecs_properties_override,
             "nodeOverrides": self.node_overrides,
             "retryStrategy": self.retry_strategy,
             "shareIdentifier": self.share_identifier,

tests/providers/amazon/aws/operators/test_batch.py

@@ -132,6 +132,7 @@ def test_init_defaults(self):
         assert batch_job.retry_strategy is None
         assert batch_job.container_overrides is None
         assert batch_job.array_properties is None
+        assert batch_job.ecs_properties_override is None
         assert batch_job.node_overrides is None
         assert batch_job.share_identifier is None
         assert batch_job.scheduling_priority_override is None
@@ -149,6 +150,7 @@ def test_template_fields_overrides(self):
             "job_queue",
             "container_overrides",
             "array_properties",
+            "ecs_properties_override",
             "node_overrides",
             "parameters",
             "retry_strategy",
@@ -204,6 +206,62 @@ def test_execute_with_failures(self):
             tags={},
         )
 
+    @mock.patch.object(BatchClientHook, "get_job_description")
+    @mock.patch.object(BatchClientHook, "wait_for_job")
+    @mock.patch.object(BatchClientHook, "check_job_success")
+    def test_execute_with_ecs_overrides(self, check_mock, wait_mock, job_description_mock):
+        self.batch.container_overrides = None
+        self.batch.ecs_properties_override = {
+            "taskProperties": [
+                {
+                    "containers": [
+                        {
+                            "command": [
+                                "string",
+                            ],
+                            "environment": [
+                                {"name": "string", "value": "string"},
+                            ],
+                            "name": "string",
+                            "resourceRequirements": [
+                                {"value": "string", "type": "'GPU'|'VCPU'|'MEMORY'"},
+                            ],
+                        },
+                    ]
+                },
+            ]
+        }
+        self.batch.execute(self.mock_context)
+
+        self.client_mock.submit_job.assert_called_once_with(
+            jobQueue="queue",
+            jobName=JOB_NAME,
+            jobDefinition="hello-world",
+            ecsPropertiesOverride={
+                "taskProperties": [
+                    {
+                        "containers": [
+                            {
+                                "command": [
+                                    "string",
+                                ],
+                                "environment": [
+                                    {"name": "string", "value": "string"},
+                                ],
+                                "name": "string",
+                                "resourceRequirements": [
+                                    {"value": "string", "type": "'GPU'|'VCPU'|'MEMORY'"},
+                                ],
+                            },
+                        ]
+                    },
+                ]
+            },
+            parameters={},
+            retryStrategy={"attempts": 1},
+            tags={},
+        )
+
     @mock.patch.object(BatchClientHook, "check_job_success")
     def test_wait_job_complete_using_waiters(self, check_mock):
         mock_waiters = mock.Mock()
@@ -238,7 +296,7 @@ def test_kill_job(self):
         self.batch.on_kill()
         self.client_mock.terminate_job.assert_called_once_with(jobId=JOB_ID, reason="Task killed by the user")
 
-    @pytest.mark.parametrize("override", ["overrides", "node_overrides"])
+    @pytest.mark.parametrize("override", ["overrides", "node_overrides", "ecs_properties_override"])
     @patch(
         "airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.client",
         new_callable=mock.PropertyMock,
@@ -269,10 +327,15 @@ def test_override_not_sent_if_not_set(self, client_mock, override):
             "parameters": {},
             "tags": {},
         }
-        if override == "overrides":
-            expected_args["containerOverrides"] = {"a": "a"}
-        else:
-            expected_args["nodeOverrides"] = {"a": "a"}
+
+        py2api = {
+            "overrides": "containerOverrides",
+            "node_overrides": "nodeOverrides",
+            "ecs_properties_override": "ecsPropertiesOverride",
+        }
+
+        expected_args[py2api[override]] = {"a": "a"}
+
         client_mock().submit_job.assert_called_once_with(**expected_args)
 
     def test_deprecated_override_param(self):