Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
camilleanne wants to merge 16 commits into from
Closed

Add support in AWS Batch Operator for multinode jobs #28321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
917c415
add node_overrides parameter to batch operator to support multinode j…
camilleanne Dec 2, 2022
57d8675
use trim_none_values to pass only truthy parameters to boto
camilleanne Dec 9, 2022
e9cf26e
add test
camilleanne Dec 9, 2022
5713c48
access logstreamname for multinode jobs; add batch_client test
camilleanne Dec 9, 2022
4c226fa
better conditionals on attempts array length
camilleanne Dec 9, 2022
d59729f
lint
camilleanne Dec 12, 2022
59e9a87
fix line length; extend test for multiple attempts
camilleanne Dec 12, 2022
5d16298
Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…
camilleanne Dec 12, 2022
c2182e3
fix bad tab
camilleanne Dec 16, 2022
830bbe9
update logstream tests
camilleanne Dec 16, 2022
b6750cf
update tests for new expectations around arrayProperties
camilleanne Dec 16, 2022
6169610
Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…
camilleanne Dec 16, 2022
79e488c
rename overrides param
vandonr-amz Dec 22, 2022
40453ca
Merge pull request #2 from aws-mwaa/vandonr/multinode
camilleanne Dec 23, 2022
179fc21
raise exception instead of a warning on unrecognized job type
camilleanne Dec 24, 2022
528bb99
Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…
camilleanne Jan 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 36 additions & 3 deletions airflow/providers/amazon/aws/hooks/batch_client.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})

@vandonr-amzvandonr-amzDec 16, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?

)
# "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.
Expand All@@ -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.
Expand Down
69 changes: 53 additions & 16 deletions airflow/providers/amazon/aws/operators/batch.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down
55 changes: 54 additions & 1 deletion tests/providers/amazon/aws/hooks/test_batch_client.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
},
}
]
}
Expand All@@ -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": [
Expand All@@ -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,
},
}
},
}
],
},
}
]
}
Comment on lines +351 to +359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄

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)
Expand Down
86 changes: 77 additions & 9 deletions tests/providers/amazon/aws/operators/test_batch.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
)
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
jobName=JOB_NAME,
containerOverrides={},
jobDefinition="hello-world",
arrayProperties={},
parameters={},
tags={},
)
Expand All@@ -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)

Expand All@@ -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")
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
Add support in AWS Batch Operator for multinode jobs#28321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

camilleanne
wants to merge
16
commits into

from

 Closed


Add support in AWS Batch Operator for multinode jobs
#28321
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


16 commits


Select commit
Hold shift + click to select a range
917c415

add node_overrides parameter to batch operator to support multinode j…


camilleanne Dec 2, 2022
57d8675

use trim_none_values to pass only truthy parameters to boto


camilleanne Dec 9, 2022
e9cf26e

add test


camilleanne Dec 9, 2022
5713c48

access logstreamname for multinode jobs; add batch_client test


camilleanne Dec 9, 2022
4c226fa

better conditionals on attempts array length


camilleanne Dec 9, 2022
d59729f

lint


camilleanne Dec 12, 2022
59e9a87

fix line length; extend test for multiple attempts


camilleanne Dec 12, 2022
5d16298

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 12, 2022
c2182e3

fix bad tab


camilleanne Dec 16, 2022
830bbe9

update logstream tests


camilleanne Dec 16, 2022
b6750cf

update tests for new expectations around arrayProperties


camilleanne Dec 16, 2022
6169610

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Dec 16, 2022
79e488c

rename overrides param


vandonr-amz Dec 22, 2022
40453ca

Merge pull request #2 from aws-mwaa/vandonr/multinode


camilleanne Dec 23, 2022
179fc21

raise exception instead of a warning on unrecognized job type


camilleanne Dec 24, 2022
528bb99

Merge branch 'main' of https://github.com/apache/airflow into ct/aws-…


camilleanne Jan 26, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


39 changes: 36 additions & 3 deletions
airflow/providers/amazon/aws/hooks/batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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", {})
@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have zero element in the array ? i.e. should we add a check on len == 0 and a user-friendly error message ?
 )
 # "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.
Expand All@@ -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.
Expand Down

69 changes: 53 additions & 16 deletions
airflow/providers/amazon/aws/operators/batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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):
Expand All@@ -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,
Expand All@@ -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 {}
Expand DownExpand Up@@ -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",
Expand Down

55 changes: 54 additions & 1 deletion
tests/providers/amazon/aws/hooks/test_batch_client.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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": {}
 },
 }
 ]
 }
Expand All@@ -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": [
Expand All@@ -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,
 },
 }
 },
 }
 ],
 },
 }
 ]
 }

Comment on lines

+351
 to 
+359

@vandonr-amzvandonr-amzDec 16, 2022

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

beautiful 😄
 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)
Expand Down

86 changes: 77 additions & 9 deletions
tests/providers/amazon/aws/operators/test_batch.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@
from __future__ import annotations

from unittest import mock
from unittest.mock import patch

import pytest

Expand DownExpand Up@@ -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",
Expand All@@ -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",
Expand DownExpand Up@@ -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
Expand All@@ -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",
Expand All@@ -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={},
 )
Expand All@@ -155,7 +157,6 @@ def test_execute_with_failures(self):
 jobName=JOB_NAME,
 containerOverrides={},
 jobDefinition="hello-world",
 arrayProperties={},
 parameters={},
 tags={},
 )
Expand All@@ -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)

Expand All@@ -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")
Expand Down