dazza-codes commented on a change in pull request #6811: [RFC][AIRFLOW-6245] Add custom waiters for AWS batch jobs URL: https://github.com/apache/airflow/pull/6811#discussion_r360731115
########## File path: airflow/providers/amazon/aws/operators/batch.py ########## @@ -19,114 +19,77 @@ # """ -Airflow operator for AWS batch service +An Airflow operator for AWS batch services -.. seealso:: http://boto3.readthedocs.io/en/latest/reference/services/batch.html -""" - -import sys -from random import randint -from time import sleep -from typing import Optional +.. seealso:: -import botocore.exceptions -import botocore.waiter + - http://boto3.readthedocs.io/en/latest/guide/configuration.html + - http://boto3.readthedocs.io/en/latest/reference/services/batch.html + - https://docs.aws.amazon.com/batch/latest/APIReference/Welcome.html +""" -from airflow.contrib.hooks.aws_hook import AwsHook from airflow.exceptions import AirflowException from airflow.models import BaseOperator -from airflow.typing_compat import Protocol +from airflow.providers.amazon.aws.hooks.batch_client import AwsBatchClient from airflow.utils.decorators import apply_defaults -# pylint: disable=invalid-name, unused-argument - -class BatchProtocol(Protocol): - """ - .. seealso:: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/batch.html - """ - - def describe_jobs(self, jobs) -> dict: - """Get job descriptions from AWS batch""" - ... - - def get_waiter(self, x: str) -> botocore.waiter.Waiter: - """Get an AWS service waiter - - Note that AWS batch might not have any waiters (until botocore PR1307 is merged and released). - - .. code-block:: python - - import boto3 - boto3.client('batch').waiter_names == [] - - .. seealso:: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html#waiters - .. seealso:: https://github.com/boto/botocore/pull/1307 - """ - ... - - def submit_job( - self, jobName, jobQueue, jobDefinition, arrayProperties, parameters, containerOverrides - ) -> dict: - """Submit a batch job - :type jobName: str - :type jobQueue: str - :type jobDefinition: str - :type arrayProperties: dict - :type parameters: dict - :type containerOverrides: dict - """ - ... - - def terminate_job(self, jobId: str, reason: str) -> dict: - """Terminate a batch job""" - ... - - -class AwsBatchOperator(BaseOperator): +class AwsBatchOperator(BaseOperator, AwsBatchClient): """ Execute a job on AWS Batch Service - .. warning: the queue parameter was renamed to job_queue to segregate the - internal CeleryExecutor queue from the AWS Batch internal queue. - :param job_name: the name for the job that will run on AWS Batch (templated) :type job_name: str + :param job_definition: the job definition name on AWS Batch :type job_definition: str + :param job_queue: the queue name on AWS Batch :type job_queue: str - :param overrides: the same parameter that boto3 will receive on - containerOverrides (templated) - http://boto3.readthedocs.io/en/latest/reference/services/batch.html#Batch.Client.submit_job + + :param overrides: the `containerOverrides` parameter for boto3 (templated) :type overrides: dict - :param array_properties: the same parameter that boto3 will receive on - arrayProperties - http://boto3.readthedocs.io/en/latest/reference/services/batch.html#Batch.Client.submit_job + + :param array_properties: the `arrayProperties` parameter for boto3 :type array_properties: dict - :param parameters: the same parameter that boto3 will receive on - parameters (templated) - http://boto3.readthedocs.io/en/latest/reference/services/batch.html#Batch.Client.submit_job + + :param parameters: the `parameters` for boto3 (templated) :type parameters: dict - :param max_retries: exponential backoff retries while waiter is not - merged, 4200 = 48 hours + + :param waiters: an :py:class:`.AwsBatchWaiters` object (see note below); + if None, polling is used with max_retries and status_retries. + :type waiters: Optional[AwsBatchWaiters] + + :param max_retries: exponential back-off retries, 4200 = 48 hours; + polling is only used when waiters is None :type max_retries: int - :param status_retries: number of retries to get job description (status), 10 + + :param status_retries: number of HTTP retries to get job status, 10; + polling is only used when waiters is None :type status_retries: int + :param aws_conn_id: connection id of AWS credentials / region name. If None, - credential boto3 strategy will be used - (http://boto3.readthedocs.io/en/latest/guide/configuration.html). + credential boto3 strategy will be used. :type aws_conn_id: str + :param region_name: region name to use in AWS Hook. Override the region_name in connection (if provided) :type region_name: str - """ - MAX_RETRIES = 4200 - STATUS_RETRIES = 10 + .. note:: + Any custom waiters must return a waiter for these calls: + .. code-block:: python + + waiter = waiters.get_waiter("JobExists") + waiter = waiters.get_waiter("JobRunning") + waiter = waiters.get_waiter("JobComplete") + + .. warning:: + The queue parameter was renamed to job_queue to segregate the + internal CeleryExecutor queue from the AWS Batch internal queue. Review comment: OK, removed the documentation and, yes, it is already noted in `UPDATING.md`. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
