Add support in AWS Batch Operator for multinode jobs

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

@@ -419,8 +419,42 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None:
 
         :param job_id: AWS Batch Job ID
         """
-        job_container_desc = self.get_job_description(job_id=job_id).get("container", {})
-        log_configuration = job_container_desc.get("logConfiguration", {})
+        job_desc = self.get_job_description(job_id=job_id)
+
+        job_node_properties = job_desc.get("nodeProperties", {})
+        job_container_desc = job_desc.get("container", {})
+
+        if job_node_properties:
+            job_node_range_properties = job_node_properties.get("nodeRangeProperties", {})
+            if len(job_node_range_properties) > 1:
+                self.log.warning(
+                    "AWS Batch job (%s) has more than one node group. Only returning logs from first group.",
+                    job_id,
+                )
+            log_configuration = (
+                job_node_range_properties[0].get("container", {}).get("logConfiguration", {})
+            )
+            # "logStreamName" value is not available in the "container" object for multinode jobs --
+            # it is available in the "attempts" object
+            job_attempts = job_desc.get("attempts", [])
+            if len(job_attempts):
+                if len(job_attempts) > 1:
+                    self.log.warning(
+                        "AWS Batch job (%s) has had more than one attempt. \
+                            Only returning logs from the most recent attempt.",
+                        job_id,
+                    )
+                awslogs_stream_name = job_attempts[-1].get("container", {}).get("logStreamName")
+            else:
+                awslogs_stream_name = None
+
+        elif job_container_desc:
+            log_configuration = job_container_desc.get("logConfiguration", {})
+            awslogs_stream_name = job_container_desc.get("logStreamName")
+        else:
+            raise AirflowException(
+                "AWS Batch job (%s) is not a supported job type. Supported job types: container, array, multinode."
+            )
 
         # In case if user select other "logDriver" rather than "awslogs"
         # than CloudWatch logging should be disabled.
@@ -435,7 +469,6 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None:
             )
             return None
 
-        awslogs_stream_name = job_container_desc.get("logStreamName")
         if not awslogs_stream_name:
             # In case of call this method on very early stage of running AWS Batch
             # there is possibility than AWS CloudWatch Stream Name not exists yet.

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

@@ -54,7 +54,13 @@ class BatchOperator(BaseOperator):
     :param job_name: the name for the job that will run on AWS Batch (templated)
     :param job_definition: the job definition name on AWS Batch
     :param job_queue: the queue name on AWS Batch
-    :param overrides: the `containerOverrides` parameter for boto3 (templated)
+
+    :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 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
@@ -88,14 +94,19 @@ class BatchOperator(BaseOperator):
         "job_name",
         "job_definition",
         "job_queue",
-        "overrides",
+        "container_overrides",
         "array_properties",
+        "node_overrides",
         "parameters",
         "waiters",
         "tags",
         "wait_for_completion",
     )
-    template_fields_renderers = {"overrides": "json", "parameters": "json"}
+    template_fields_renderers = {
+        "container_overrides": "json",
+        "parameters": "json",
+        "node_overrides": "json",
+    }
 
     @property
     def operator_extra_links(self):
@@ -114,8 +125,10 @@ def __init__(
         job_name: str,
         job_definition: str,
         job_queue: str,
-        overrides: dict,
+        overrides: dict | None = None,  # deprecated
+        container_overrides: dict | None = None,
         array_properties: dict | None = None,
+        node_overrides: dict | None = None,
         parameters: dict | None = None,
         job_id: str | None = None,
         waiters: Any | None = None,
@@ -133,8 +146,23 @@ def __init__(
         self.job_name = job_name
         self.job_definition = job_definition
         self.job_queue = job_queue
-        self.overrides = overrides or {}
-        self.array_properties = array_properties or {}
+
+        if overrides:
+            self.container_overrides = overrides
+            warnings.warn(
+                f"Parameter `overrides` is deprecated, Please use `container_overrides` instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if container_overrides:
+                raise AirflowException(
+                    "If providing `container_overrides`, then old parameter 'overrides' should be removed."
+                )
+        else:
+            self.container_overrides = container_overrides
+
+        self.node_overrides = node_overrides
+        self.array_properties = array_properties
         self.parameters = parameters or {}
         self.waiters = waiters
         self.tags = tags or {}
@@ -174,18 +202,27 @@ def submit_job(self, context: Context):
             self.job_definition,
             self.job_queue,
         )
-        self.log.info("AWS Batch job - container overrides: %s", self.overrides)
+
+        if self.container_overrides:
+            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.node_overrides:
+            self.log.info("AWS Batch job - node properties: %s", self.node_overrides)
+
+        args = {
+            "jobName": self.job_name,
+            "jobQueue": self.job_queue,
+            "jobDefinition": self.job_definition,
+            "arrayProperties": self.array_properties,
+            "parameters": self.parameters,
+            "tags": self.tags,
+            "containerOverrides": self.container_overrides,
+            "nodeOverrides": self.node_overrides,
+        }
 
         try:
-            response = self.hook.client.submit_job(
-                jobName=self.job_name,
-                jobQueue=self.job_queue,
-                jobDefinition=self.job_definition,
-                arrayProperties=self.array_properties,
-                parameters=self.parameters,
-                containerOverrides=self.overrides,
-                tags=self.tags,
-            )
+            response = self.hook.client.submit_job(**trim_none_values(args))
         except Exception as e:
             self.log.error(
                 "AWS Batch job failed submission - job definition: %s - on queue %s",

tests/providers/amazon/aws/hooks/test_batch_client.py

@@ -274,13 +274,16 @@ def test_job_awslogs_user_defined(self):
         assert awslogs["awslogs_stream_name"] == LOG_STREAM_NAME
         assert awslogs["awslogs_group"] == "/test/batch/job"
         assert awslogs["awslogs_region"] == "ap-southeast-2"
+
 
     def test_job_no_awslogs_stream(self, caplog):
         self.client_mock.describe_jobs.return_value = {
             "jobs": [
                 {
                     "jobId": JOB_ID,
-                    "container": {},
+                    "container": {
+                        "logConfiguration": {}
+                    },
                 }
             ]
         }
@@ -290,6 +293,22 @@ def test_job_no_awslogs_stream(self, caplog):
             assert len(caplog.records) == 1
             assert "doesn't create AWS CloudWatch Stream" in caplog.messages[0]
 
+    def test_job_not_recognized_job(self):
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID
+                }
+            ]
+        }
+        with pytest.raises(AirflowException) as ctx:
+            self.batch_client.get_job_awslogs_info(JOB_ID)
+        # It should not retry when this client error occurs
+        self.client_mock.describe_jobs.assert_called_once_with(jobs=[JOB_ID])
+        msg = f"AWS Batch job (%s) is not a supported job type. Supported job types: container, array, multinode."
+        assert msg in str(ctx.value)
+
+
     def test_job_splunk_logs(self, caplog):
         self.client_mock.describe_jobs.return_value = {
             "jobs": [
@@ -309,6 +328,40 @@ def test_job_splunk_logs(self, caplog):
             assert len(caplog.records) == 1
             assert "uses logDriver (splunk). AWS CloudWatch logging disabled." in caplog.messages[0]
 
+    def test_job_awslogs_multinode_job(self):
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID,
+                    "attempts": [
+                        {"container": {"exitCode": 0, "logStreamName": "test/stream/attempt0"}},
+                        {"container": {"exitCode": 0, "logStreamName": LOG_STREAM_NAME}},
+                    ],
+                    "nodeProperties": {
+                        "mainNode": 0,
+                        "nodeRangeProperties": [
+                            {
+                                "targetNodes": "0:",
+                                "container": {
+                                    "logConfiguration": {
+                                        "logDriver": "awslogs",
+                                        "options": {
+                                            "awslogs-group": "/test/batch/job",
+                                            "awslogs-region": AWS_REGION,
+                                        },
+                                    }
+                                },
+                            }
+                        ],
+                    },
+                }
+            ]
+        }
+        awslogs = self.batch_client.get_job_awslogs_info(JOB_ID)
+        assert awslogs["awslogs_stream_name"] == LOG_STREAM_NAME
+        assert awslogs["awslogs_group"] == "/test/batch/job"
+        assert awslogs["awslogs_region"] == AWS_REGION
+
 
 class TestBatchClientDelays:
     @mock.patch.dict("os.environ", AWS_DEFAULT_REGION=AWS_REGION)

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

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 from unittest import mock
+from unittest.mock import patch
 
 import pytest
 
@@ -48,7 +49,7 @@ class TestBatchOperator:
     @mock.patch.dict("os.environ", AWS_ACCESS_KEY_ID=AWS_ACCESS_KEY_ID)
     @mock.patch.dict("os.environ", AWS_SECRET_ACCESS_KEY=AWS_SECRET_ACCESS_KEY)
     @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.AwsBaseHook.get_client_type")
-    def setup_method(self, method, get_client_type_mock):
+    def setup_method(self, _, get_client_type_mock):
         self.get_client_type_mock = get_client_type_mock
         self.batch = BatchOperator(
             task_id="task",
@@ -58,7 +59,7 @@ def setup_method(self, method, get_client_type_mock):
             max_retries=self.MAX_RETRIES,
             status_retries=self.STATUS_RETRIES,
             parameters=None,
-            overrides={},
+            container_overrides={},
             array_properties=None,
             aws_conn_id="airflow_test",
             region_name="eu-west-1",
@@ -91,8 +92,9 @@ def test_init(self):
         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.overrides == {}
-        assert self.batch.array_properties == {}
+        assert self.batch.container_overrides == {}
+        assert self.batch.array_properties is None
+        assert self.batch.node_overrides 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
@@ -107,8 +109,9 @@ def test_template_fields_overrides(self):
             "job_name",
             "job_definition",
             "job_queue",
-            "overrides",
+            "container_overrides",
             "array_properties",
+            "node_overrides",
             "parameters",
             "waiters",
             "tags",
@@ -131,7 +134,6 @@ def test_execute_without_failures(self, check_mock, wait_mock, job_description_m
             jobName=JOB_NAME,
             containerOverrides={},
             jobDefinition="hello-world",
-            arrayProperties={},
             parameters={},
             tags={},
         )
@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
             jobName=JOB_NAME,
             containerOverrides={},
             jobDefinition="hello-world",
-            arrayProperties={},
             parameters={},
             tags={},
         )
@@ -166,9 +167,17 @@ def test_wait_job_complete_using_waiters(self, check_mock):
         self.batch.waiters = mock_waiters
 
         self.client_mock.submit_job.return_value = RESPONSE_WITHOUT_FAILURES
-        self.client_mock.describe_jobs.return_value = {"jobs": [{"jobId": JOB_ID, "status": "SUCCEEDED"}]}
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID,
+                    "status": "SUCCEEDED",
+                    "logStreamName": "logStreamName",
+                    "container": {"logConfiguration": {}},
+                }
+            ]
+        }
         self.batch.execute(self.mock_context)
-
         mock_waiters.wait_for_job.assert_called_once_with(JOB_ID)
         check_mock.assert_called_once_with(JOB_ID)
 
@@ -186,6 +195,65 @@ 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"])
+    @patch(
+        "airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.client",
+        new_callable=mock.PropertyMock,
+    )
+    def test_override_not_sent_if_not_set(self, client_mock, override):
+        """
+        check that when setting container override or node override, the other key is not sent
+        in the API call (which would create a validation error from boto)
+        """
+        override_arg = {override: {"a": "a"}}
+        batch = BatchOperator(
+            task_id="task",
+            job_name=JOB_NAME,
+            job_queue="queue",
+            job_definition="hello-world",
+            **override_arg,
+            # setting those to bypass code that is not relevant here
+            do_xcom_push=False,
+            wait_for_completion=False,
+        )
+
+        batch.execute(None)
+
+        expected_args = {
+            "jobQueue": "queue",
+            "jobName": JOB_NAME,
+            "jobDefinition": "hello-world",
+            "parameters": {},
+            "tags": {},
+        }
+        if override == "overrides":
+            expected_args["containerOverrides"] = {"a": "a"}
+        else:
+            expected_args["nodeOverrides"] = {"a": "a"}
+        client_mock().submit_job.assert_called_once_with(**expected_args)
+
+    def test_deprecated_override_param(self):
+        with pytest.warns(DeprecationWarning):
+            _ = BatchOperator(
+                task_id="task",
+                job_name=JOB_NAME,
+                job_queue="queue",
+                job_definition="hello-world",
+                overrides={"a": "b"},  # <- the deprecated field
+            )
+
+    def test_cant_set_old_and_new_override_param(self):
+        with pytest.raises(AirflowException):
+            _ = BatchOperator(
+                task_id="task",
+                job_name=JOB_NAME,
+                job_queue="queue",
+                job_definition="hello-world",
+                # can't set both of those, as one is a replacement for the other
+                overrides={"a": "b"},
+                container_overrides={"a": "b"},
+            )
+
 
 class TestBatchCreateComputeEnvironmentOperator:
     @mock.patch.object(BatchClientHook, "client")