From 917c41521aa764060b13745d52ffb379341ab8d6 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Fri, 2 Dec 2022 15:29:17 -0800 Subject: [PATCH 01/12] add node_overrides parameter to batch operator to support multinode jobs; update client to collect log info from multinode job descriptions --- .../amazon/aws/hooks/batch_client.py | 20 ++++++++++++++-- .../providers/amazon/aws/operators/batch.py | 24 ++++++++++++++----- 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index e9080189e76b2..a0c563f5e0015 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -412,8 +412,24 @@ 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", {}) + elif job_container_desc: + log_configuration = job_container_desc.get("logConfiguration", {}) + else: + self.log.warning("AWS Batch job (%s) is neither a container nor multinode job. Log info not found.") + return None # In case if user select other "logDriver" rather than "awslogs" # than CloudWatch logging should be disabled. diff --git a/airflow/providers/amazon/aws/operators/batch.py b/airflow/providers/amazon/aws/operators/batch.py index 9e85afaf4c0ce..aeffde462b5d5 100644 --- a/airflow/providers/amazon/aws/operators/batch.py +++ b/airflow/providers/amazon/aws/operators/batch.py @@ -65,6 +65,8 @@ class BatchOperator(BaseOperator): :param 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) @@ -108,12 +110,13 @@ class BatchOperator(BaseOperator): "job_queue", "overrides", "array_properties", + "node_overrides", "parameters", "waiters", "tags", "wait_for_completion", ) - template_fields_renderers = {"overrides": "json", "parameters": "json"} + template_fields_renderers = {"overrides": "json", "parameters": "json", "node_overrides": "json"} @property def operator_extra_links(self): @@ -132,8 +135,9 @@ def __init__( job_name: str, job_definition: str, job_queue: str, - overrides: dict, + 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, @@ -151,8 +155,9 @@ 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 {} + self.container_overrides = overrides or None + self.array_properties = array_properties or None + self.node_overrides = node_overrides or None self.parameters = parameters or {} self.waiters = waiters self.tags = tags or {} @@ -192,7 +197,13 @@ 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) try: response = self.hook.client.submit_job( @@ -201,7 +212,8 @@ def submit_job(self, context: Context): jobDefinition=self.job_definition, arrayProperties=self.array_properties, parameters=self.parameters, - containerOverrides=self.overrides, + containerOverrides=self.container_overrides, + nodeOverrides=self.node_overrides, tags=self.tags, ) except Exception as e: From 57d8675002ab150349e515bb93e77aabe6f0e723 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Thu, 8 Dec 2022 17:25:15 -0800 Subject: [PATCH 02/12] use trim_none_values to pass only truthy parameters to boto --- .../providers/amazon/aws/operators/batch.py | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/batch.py b/airflow/providers/amazon/aws/operators/batch.py index aeffde462b5d5..8476185c82e42 100644 --- a/airflow/providers/amazon/aws/operators/batch.py +++ b/airflow/providers/amazon/aws/operators/batch.py @@ -155,9 +155,9 @@ def __init__( self.job_name = job_name self.job_definition = job_definition self.job_queue = job_queue - self.container_overrides = overrides or None - self.array_properties = array_properties or None - self.node_overrides = node_overrides or None + self.container_overrides = overrides + self.node_overrides = node_overrides + self.array_properties = array_properties self.parameters = parameters or {} self.waiters = waiters self.tags = tags or {} @@ -205,17 +205,20 @@ def submit_job(self, context: Context): 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.container_overrides, - nodeOverrides=self.node_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", From e9cf26e700363fc62bb77269fa8a7e58e973461e Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Thu, 8 Dec 2022 17:33:47 -0800 Subject: [PATCH 03/12] add test --- .../amazon/aws/operators/test_batch.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index c952680ba14c9..131318ae21691 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -19,6 +19,7 @@ import unittest from unittest import mock +from unittest.mock import patch import pytest @@ -109,6 +110,7 @@ def test_template_fields_overrides(self): "job_definition", "job_queue", "overrides", + "node_overrides", "array_properties", "parameters", "waiters", @@ -187,6 +189,45 @@ 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") +class TestBatchOperator2: + """test class that does not inherit from unittest.TestCase""" + @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", + "arrayProperties": {}, + "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) + class TestBatchCreateComputeEnvironmentOperator(unittest.TestCase): @mock.patch.object(BatchClientHook, "client") From 5713c484d3f1504b0498fa29f0def510cf97c709 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Fri, 9 Dec 2022 14:48:51 -0800 Subject: [PATCH 04/12] access logstreamname for multinode jobs; add batch_client test --- .../amazon/aws/hooks/batch_client.py | 14 ++++++- .../amazon/aws/hooks/test_batch_client.py | 37 +++++++++++++++++++ .../amazon/aws/operators/test_batch.py | 2 +- 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index a0c563f5e0015..2d85deee66566 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -425,8 +425,21 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: 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) > 1: + self.log.warning( + "AWS Batch job (%s) has had more than one attempt. Only returning logs from the most recent attempt.", + job_id, + ) + elif not len(job_attempts): + awslogs_stream_name = None + awslogs_stream_name = job_attempts[-1].get("container", {}).get("logStreamName") + elif job_container_desc: log_configuration = job_container_desc.get("logConfiguration", {}) + awslogs_stream_name = job_container_desc.get("logStreamName") else: self.log.warning("AWS Batch job (%s) is neither a container nor multinode job. Log info not found.") return None @@ -444,7 +457,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. diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py index 030406c0cde72..cf4db5e58ab3a 100644 --- a/tests/providers/amazon/aws/hooks/test_batch_client.py +++ b/tests/providers/amazon/aws/hooks/test_batch_client.py @@ -303,6 +303,43 @@ def test_job_splunk_logs(self): assert self.batch_client.get_job_awslogs_info(JOB_ID) is None assert len(capture_logs.records) == 1 + def test_job_awslogs_multinode_job(self): + self.client_mock.describe_jobs.return_value = { + "jobs": [ + { + "jobId": JOB_ID, + "attempts": [ + { + "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(unittest.TestCase): @mock.patch.dict("os.environ", AWS_DEFAULT_REGION=AWS_REGION) diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index 131318ae21691..aa9373e77a45c 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -189,7 +189,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") -class TestBatchOperator2: +class TestBatchOperatorTrimmedArgs: """test class that does not inherit from unittest.TestCase""" @pytest.mark.parametrize("override", ["overrides", "node_overrides"]) @patch("airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.client", From 4c226fa73f7be3321fffb489ab14c992c42484ac Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Fri, 9 Dec 2022 14:52:10 -0800 Subject: [PATCH 05/12] better conditionals on attempts array length --- .../providers/amazon/aws/hooks/batch_client.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index 2d85deee66566..c171e08fd6e3a 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -428,14 +428,15 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: # "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) > 1: - self.log.warning( - "AWS Batch job (%s) has had more than one attempt. Only returning logs from the most recent attempt.", - job_id, - ) - elif not len(job_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 - awslogs_stream_name = job_attempts[-1].get("container", {}).get("logStreamName") elif job_container_desc: log_configuration = job_container_desc.get("logConfiguration", {}) From d59729f4f6eb0727b0a8082f6105086ca7390405 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Mon, 12 Dec 2022 10:12:22 -0800 Subject: [PATCH 06/12] lint --- .../providers/amazon/aws/hooks/batch_client.py | 8 ++++++-- .../providers/amazon/aws/operators/batch.py | 7 +++---- .../amazon/aws/hooks/test_batch_client.py | 18 ++++++------------ .../amazon/aws/operators/test_batch.py | 9 ++++++--- 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index c171e08fd6e3a..cc645c67c26b8 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -424,7 +424,9 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: "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", {}) + 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", []) @@ -442,7 +444,9 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: log_configuration = job_container_desc.get("logConfiguration", {}) awslogs_stream_name = job_container_desc.get("logStreamName") else: - self.log.warning("AWS Batch job (%s) is neither a container nor multinode job. Log info not found.") + self.log.warning( + "AWS Batch job (%s) is neither a container nor multinode job. Log info not found." + ) return None # In case if user select other "logDriver" rather than "awslogs" diff --git a/airflow/providers/amazon/aws/operators/batch.py b/airflow/providers/amazon/aws/operators/batch.py index 8476185c82e42..4035a788be381 100644 --- a/airflow/providers/amazon/aws/operators/batch.py +++ b/airflow/providers/amazon/aws/operators/batch.py @@ -198,14 +198,13 @@ def submit_job(self, context: Context): self.job_queue, ) - if (self.container_overrides): + if self.container_overrides: self.log.info("AWS Batch job - container overrides: %s", self.container_overrides) - if (self.array_properties): + if self.array_properties: self.log.info("AWS Batch job - array properties: %s", self.array_properties) - if (self.node_overrides): + 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, diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py index cf4db5e58ab3a..85f3cae2bce48 100644 --- a/tests/providers/amazon/aws/hooks/test_batch_client.py +++ b/tests/providers/amazon/aws/hooks/test_batch_client.py @@ -308,14 +308,7 @@ def test_job_awslogs_multinode_job(self): "jobs": [ { "jobId": JOB_ID, - "attempts": [ - { - "container": { - "exitCode": 0, - "logStreamName": LOG_STREAM_NAME - } - } - ], + "attempts": [{"container": {"exitCode": 0, "logStreamName": LOG_STREAM_NAME}}], "nodeProperties": { "mainNode": 0, "nodeRangeProperties": [ @@ -326,13 +319,13 @@ def test_job_awslogs_multinode_job(self): "logDriver": "awslogs", "options": { "awslogs-group": "/test/batch/job", - "awslogs-region": AWS_REGION + "awslogs-region": AWS_REGION, }, } - } + }, } - ] - } + ], + }, } ] } @@ -341,6 +334,7 @@ def test_job_awslogs_multinode_job(self): assert awslogs["awslogs_group"] == "/test/batch/job" assert awslogs["awslogs_region"] == AWS_REGION + class TestBatchClientDelays(unittest.TestCase): @mock.patch.dict("os.environ", AWS_DEFAULT_REGION=AWS_REGION) @mock.patch.dict("os.environ", AWS_ACCESS_KEY_ID=AWS_ACCESS_KEY_ID) diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index aa9373e77a45c..c5c813279a39a 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -189,11 +189,15 @@ 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") + class TestBatchOperatorTrimmedArgs: """test class that does not inherit from unittest.TestCase""" + @pytest.mark.parametrize("override", ["overrides", "node_overrides"]) - @patch("airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.client", - new_callable=mock.PropertyMock) + @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 @@ -206,7 +210,6 @@ def test_override_not_sent_if_not_set(self, client_mock, override): 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, From 59e9a870c8a10b0e919aac6856aad303ef434f5e Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Mon, 12 Dec 2022 15:07:47 -0800 Subject: [PATCH 07/12] fix line length; extend test for multiple attempts --- airflow/providers/amazon/aws/hooks/batch_client.py | 3 ++- tests/providers/amazon/aws/hooks/test_batch_client.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index cc645c67c26b8..c81b729e4eb02 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -433,7 +433,8 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: 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.", + "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") diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py index 85f3cae2bce48..83129f247d2ef 100644 --- a/tests/providers/amazon/aws/hooks/test_batch_client.py +++ b/tests/providers/amazon/aws/hooks/test_batch_client.py @@ -308,7 +308,10 @@ def test_job_awslogs_multinode_job(self): "jobs": [ { "jobId": JOB_ID, - "attempts": [{"container": {"exitCode": 0, "logStreamName": LOG_STREAM_NAME}}], + "attempts": [ + {"container": {"exitCode": 0, "logStreamName": "test/stream/attempt0"}}, + {"container": {"exitCode": 0, "logStreamName": LOG_STREAM_NAME}}, + ], "nodeProperties": { "mainNode": 0, "nodeRangeProperties": [ From c2182e3be549b8fdaf6c488bf66fb1fe6e1f43df Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Thu, 15 Dec 2022 16:31:03 -0800 Subject: [PATCH 08/12] fix bad tab --- .../amazon/aws/hooks/batch_client.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index 7c57f70d6b50c..b02c7e1c06d53 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -428,22 +428,22 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: "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 + 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", {}) From 830bbe9cd82f6409eec9d4b250823ad3e68bb960 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Thu, 15 Dec 2022 16:31:27 -0800 Subject: [PATCH 09/12] update logstream tests --- .../amazon/aws/hooks/test_batch_client.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py index fcc4bc7d68296..a17361e58958c 100644 --- a/tests/providers/amazon/aws/hooks/test_batch_client.py +++ b/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,20 @@ 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, caplog): + self.client_mock.describe_jobs.return_value = { + "jobs": [ + { + "jobId": JOB_ID + } + ] + } + + with caplog.at_level(level=logging.WARNING): + assert self.batch_client.get_job_awslogs_info(JOB_ID) is None + assert len(caplog.records) == 1 + assert "neither a container nor multinode job." in caplog.messages[0] + def test_job_splunk_logs(self, caplog): self.client_mock.describe_jobs.return_value = { "jobs": [ From b6750cf819cb7c0c5e41dc19a3fcb346ade7f004 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Thu, 15 Dec 2022 17:07:52 -0800 Subject: [PATCH 10/12] update tests for new expectations around arrayProperties --- tests/providers/amazon/aws/operators/test_batch.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index d80d1dfa6afc6..a21ae49531e50 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -92,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 == None + assert self.batch.node_overrides == 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 @@ -109,8 +110,8 @@ def test_template_fields_overrides(self): "job_definition", "job_queue", "overrides", - "node_overrides", "array_properties", + "node_overrides", "parameters", "waiters", "tags", @@ -133,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={}, ) @@ -157,7 +157,6 @@ def test_execute_with_failures(self): jobName=JOB_NAME, containerOverrides={}, jobDefinition="hello-world", - arrayProperties={}, parameters={}, tags={}, ) @@ -220,7 +219,6 @@ def test_override_not_sent_if_not_set(self, client_mock, override): "jobQueue": "queue", "jobName": JOB_NAME, "jobDefinition": "hello-world", - "arrayProperties": {}, "parameters": {}, "tags": {}, } From 79e488c4f1229806a8d942d2050411b1c85cc861 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Thu, 22 Dec 2022 13:10:43 -0800 Subject: [PATCH 11/12] rename overrides param --- .../providers/amazon/aws/operators/batch.py | 32 ++++++++++++++--- .../amazon/aws/operators/test_batch.py | 36 ++++++++++++++----- 2 files changed, 54 insertions(+), 14 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/batch.py b/airflow/providers/amazon/aws/operators/batch.py index 4035a788be381..6409faf1674c3 100644 --- a/airflow/providers/amazon/aws/operators/batch.py +++ b/airflow/providers/amazon/aws/operators/batch.py @@ -26,6 +26,7 @@ from __future__ import annotations import sys +import warnings from typing import TYPE_CHECKING, Any, Sequence from airflow.providers.amazon.aws.utils import trim_none_values @@ -63,7 +64,9 @@ class BatchOperator(BaseOperator): :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) @@ -108,7 +111,7 @@ class BatchOperator(BaseOperator): "job_name", "job_definition", "job_queue", - "overrides", + "container_overrides", "array_properties", "node_overrides", "parameters", @@ -116,7 +119,11 @@ class BatchOperator(BaseOperator): "tags", "wait_for_completion", ) - template_fields_renderers = {"overrides": "json", "parameters": "json", "node_overrides": "json"} + template_fields_renderers = { + "container_overrides": "json", + "parameters": "json", + "node_overrides": "json", + } @property def operator_extra_links(self): @@ -135,7 +142,8 @@ def __init__( job_name: str, job_definition: str, job_queue: str, - overrides: dict | None = None, + overrides: dict | None = None, # deprecated + container_overrides: dict | None = None, array_properties: dict | None = None, node_overrides: dict | None = None, parameters: dict | None = None, @@ -155,7 +163,21 @@ def __init__( self.job_name = job_name self.job_definition = job_definition self.job_queue = job_queue - self.container_overrides = overrides + + 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 {} diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index a21ae49531e50..dcc3943a6e163 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -49,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", @@ -59,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", @@ -93,8 +93,8 @@ def test_init(self): assert self.batch.hook.status_retries == self.STATUS_RETRIES assert self.batch.parameters == {} assert self.batch.container_overrides == {} - assert self.batch.array_properties == None - assert self.batch.node_overrides == None + 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 @@ -109,7 +109,7 @@ def test_template_fields_overrides(self): "job_name", "job_definition", "job_queue", - "overrides", + "container_overrides", "array_properties", "node_overrides", "parameters", @@ -187,10 +187,6 @@ 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") - -class TestBatchOperatorTrimmedArgs: - """test class that does not inherit from unittest.TestCase""" - @pytest.mark.parametrize("override", ["overrides", "node_overrides"]) @patch( "airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.client", @@ -228,6 +224,28 @@ def test_override_not_sent_if_not_set(self, client_mock, override): 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") From 179fc21e00a26c84bd9fba23ad9898b905959151 Mon Sep 17 00:00:00 2001 From: Camille Teicheira Date: Fri, 23 Dec 2022 16:02:16 -0800 Subject: [PATCH 12/12] raise exception instead of a warning on unrecognized job type --- airflow/providers/amazon/aws/hooks/batch_client.py | 5 ++--- .../providers/amazon/aws/hooks/test_batch_client.py | 12 +++++++----- tests/providers/amazon/aws/operators/test_batch.py | 12 ++++++++++-- 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py index b02c7e1c06d53..c06811f5ab2ca 100644 --- a/airflow/providers/amazon/aws/hooks/batch_client.py +++ b/airflow/providers/amazon/aws/hooks/batch_client.py @@ -449,10 +449,9 @@ def get_job_awslogs_info(self, job_id: str) -> dict[str, str] | None: log_configuration = job_container_desc.get("logConfiguration", {}) awslogs_stream_name = job_container_desc.get("logStreamName") else: - self.log.warning( - "AWS Batch job (%s) is neither a container nor multinode job. Log info not found." + raise AirflowException( + "AWS Batch job (%s) is not a supported job type. Supported job types: container, array, multinode." ) - return None # In case if user select other "logDriver" rather than "awslogs" # than CloudWatch logging should be disabled. diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py index a17361e58958c..86ca4aad64e55 100644 --- a/tests/providers/amazon/aws/hooks/test_batch_client.py +++ b/tests/providers/amazon/aws/hooks/test_batch_client.py @@ -293,7 +293,7 @@ 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, caplog): + def test_job_not_recognized_job(self): self.client_mock.describe_jobs.return_value = { "jobs": [ { @@ -301,11 +301,13 @@ def test_job_not_recognized_job(self, caplog): } ] } + 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) - with caplog.at_level(level=logging.WARNING): - assert self.batch_client.get_job_awslogs_info(JOB_ID) is None - assert len(caplog.records) == 1 - assert "neither a container nor multinode job." in caplog.messages[0] def test_job_splunk_logs(self, caplog): self.client_mock.describe_jobs.return_value = { diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index dcc3943a6e163..c6a923b51ec2d 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -167,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)