<s> from builtins import object <EOL> import logging <EOL> import subprocess <EOL> import time <EOL> from celery import Celery <EOL> from celery import states as celery_states <EOL> from airflow . exceptions import AirflowException <EOL> from airflow . executors . base_executor import BaseExecutor <EOL> from airflow import configuration <EOL> PARALLELISM = configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> '''<STR_LIT>''' <EOL> DEFAULT_QUEUE = configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> class CeleryConfig ( object ) : <EOL> CELERY_ACCEPT_CONTENT = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> CELERYD_PREFETCH_MULTIPLIER = <NUM_LIT:1> <EOL> CELERY_ACKS_LATE = True <EOL> BROKER_URL = configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> CELERY_RESULT_BACKEND = configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> CELERYD_CONCURRENCY = configuration . getint ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> CELERY_DEFAULT_QUEUE = DEFAULT_QUEUE <EOL> CELERY_DEFAULT_EXCHANGE = DEFAULT_QUEUE <EOL> app = Celery ( <EOL> configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> config_source = CeleryConfig ) <EOL> @ app . task <EOL> def execute_command ( command ) : <EOL> logging . info ( "<STR_LIT>" + command ) <EOL> try : <EOL> subprocess . check_call ( command , shell = True ) <EOL> except subprocess . CalledProcessError as e : <EOL> logging . error ( e ) <EOL> raise AirflowException ( '<STR_LIT>' ) <EOL> class CeleryExecutor ( BaseExecutor ) : <EOL> """<STR_LIT>""" <EOL> def start ( self ) : <EOL> self . tasks = { } <EOL> self . last_state = { } <EOL> def execute_async ( self , key , command , queue = DEFAULT_QUEUE ) : <EOL> self . logger . info ( "<STR_LIT>" <EOL> "<STR_LIT>" . format ( ** locals ( ) ) ) <EOL> self . tasks [ key ] = execute_command . apply_async ( <EOL> args = [ command ] , queue = queue ) <EOL> self . last_state [ key ] = celery_states . PENDING <EOL> def sync ( self ) : <EOL> self . logger . debug ( <EOL> "<STR_LIT>" . format ( len ( self . tasks ) ) ) <EOL> for key , async in list ( self . tasks . items ( ) ) : <EOL> state = async . state <EOL> if self . last_state [ key ] != state : <EOL> if state == celery_states . SUCCESS : <EOL> self . success ( key ) <EOL> del self . tasks [ key ] <EOL> del self . last_state [ key ] <EOL> elif state == celery_states . FAILURE : <EOL> self . fail ( key ) <EOL> del self . tasks [ key ] <EOL> del self . last_state [ key ] <EOL> elif state == celery_states . REVOKED : <EOL> self . fail ( key ) <EOL> del self . tasks [ key ] <EOL> del self . last_state [ key ] <EOL> else : <EOL> self . logger . info ( "<STR_LIT>" + async . state ) <EOL> self . last_state [ key ] = async . state <EOL> def end ( self , synchronous = False ) : <EOL> if synchronous : <EOL> while any ( [ <EOL> async . state not in celery_states . READY_STATES <EOL> for async in self . tasks . values ( ) ] ) : <EOL> time . sleep ( <NUM_LIT:5> ) </s>
