airflow
airflow copied to clipboard
Support AWS Batch multinode job types
Description
Support multinode job types in the AWS Batch Operator.
The boto3 submit_job
method supports container, multinode, and array batch jobs with the mutually exclusive nodeOverrides
and containerOverrides
(+ arrayProperties
) parameters. But currently the AWS Batch Operator only supports submission of container jobs and array jobs by hardcoding the boto3 submit_job
parameter containerOverrides
: https://github.com/apache/airflow/blob/3c08cefdfd2e2636a714bb835902f0cb34225563/airflow/providers/amazon/aws/operators/batch.py#L200 & https://github.com/apache/airflow/blob/3c08cefdfd2e2636a714bb835902f0cb34225563/airflow/providers/amazon/aws/hooks/batch_client.py#L99
The get_job_awslogs_info
method in the batch client hook is also hardcoded for the container type job: https://github.com/apache/airflow/blob/3c08cefdfd2e2636a714bb835902f0cb34225563/airflow/providers/amazon/aws/hooks/batch_client.py#L425
To support multinode jobs the get_job_awslogs_info
method would need to access nodeProperties
from the describe_jobs
response.
Use case/motivation
Multinode jobs are a supported job type of AWS Batch, are supported by the underlying boto3 library, and should be also be available to be managed by Airflow. I've extended the AWS Batch Operator for our own use cases, but would prefer to not maintain a separate operator.
Related issues
No response
Are you willing to submit a PR?
- [X] Yes I am willing to submit a PR!
Code of Conduct
- [X] I agree to follow this project's Code of Conduct
Thanks for opening your first issue here! Be sure to follow the issue template!
Feel free. I assigned you. You might also sync this up with the work done by @ferruzzi @o-nikolas and @vincbeck from the AWS team -
@camilleanne are you working on this issue?
Unassigning Camille, so that others are free to take this task, since we have not heard back.
@camilleanne has put out a PR (thanks Camille!) reassigning