airflow icon indicating copy to clipboard operation
airflow copied to clipboard

Support AWS Batch multinode job types

Open camilleanne opened this issue 2 years ago • 2 comments

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

camilleanne avatar Aug 03 '22 23:08 camilleanne

Thanks for opening your first issue here! Be sure to follow the issue template!

boring-cyborg[bot] avatar Aug 03 '22 23:08 boring-cyborg[bot]

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 -

potiuk avatar Aug 04 '22 07:08 potiuk

@camilleanne are you working on this issue?

eladkal avatar Sep 30 '22 20:09 eladkal

Unassigning Camille, so that others are free to take this task, since we have not heard back.

o-nikolas avatar Oct 17 '22 22:10 o-nikolas

@camilleanne has put out a PR (thanks Camille!) reassigning

o-nikolas avatar Dec 12 '22 23:12 o-nikolas