Coverage for airflow/executors/celery_executor.py : 45%

# -*- coding: utf-8 -*- 

# 

# Licensed under the Apache License, Version 2.0 (the "License"); 

# you may not use this file except in compliance with the License. 

# You may obtain a copy of the License at 

# 

# http://www.apache.org/licenses/LICENSE-2.0 

# 

# Unless required by applicable law or agreed to in writing, software 

# distributed under the License is distributed on an "AS IS" BASIS, 

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

# See the License for the specific language governing permissions and 

# limitations under the License. 

from builtins import object 

import logging 

import subprocess 

import time 

from celery import Celery 

from celery import states as celery_states 

from airflow.exceptions import AirflowException 

from airflow.executors.base_executor import BaseExecutor 

from airflow import configuration 

PARALLELISM = configuration.get('core', 'PARALLELISM') 

''' 

To start the celery worker, run the command: 

airflow worker 

''' 

DEFAULT_QUEUE = configuration.get('celery', 'DEFAULT_QUEUE') 

class CeleryConfig(object): 

CELERY_ACCEPT_CONTENT = ['json', 'pickle'] 

CELERYD_PREFETCH_MULTIPLIER = 1 

CELERY_ACKS_LATE = True 

BROKER_URL = configuration.get('celery', 'BROKER_URL') 

CELERY_RESULT_BACKEND = configuration.get('celery', 'CELERY_RESULT_BACKEND') 

CELERYD_CONCURRENCY = configuration.getint('celery', 'CELERYD_CONCURRENCY') 

CELERY_DEFAULT_QUEUE = DEFAULT_QUEUE 

CELERY_DEFAULT_EXCHANGE = DEFAULT_QUEUE 

app = Celery( 

configuration.get('celery', 'CELERY_APP_NAME'), 

config_source=CeleryConfig) 

@app.task 

def execute_command(command): 

logging.info("Executing command in Celery " + command) 

try: 

subprocess.check_call(command, shell=True) 

except subprocess.CalledProcessError as e: 

logging.error(e) 

raise AirflowException('Celery command failed') 

class CeleryExecutor(BaseExecutor): 

""" 

CeleryExecutor is recommended for production use of Airflow. It allows 

distributing the execution of task instances to multiple worker nodes. 

Celery is a simple, flexible and reliable distributed system to process 

vast amounts of messages, while providing operations with the tools 

required to maintain such a system. 

""" 

def start(self): 

self.tasks = {} 

self.last_state = {} 

def execute_async(self, key, command, queue=DEFAULT_QUEUE): 

self.logger.info( "[celery] queuing {key} through celery, " 

"queue={queue}".format(**locals())) 

self.tasks[key] = execute_command.apply_async( 

args=[command], queue=queue) 

self.last_state[key] = celery_states.PENDING 

def sync(self): 

self.logger.debug( 

"Inquiring about {} celery task(s)".format(len(self.tasks))) 

for key, async in list(self.tasks.items()): 

state = async.state 

if self.last_state[key] != state: 

if state == celery_states.SUCCESS: 

self.success(key) 

del self.tasks[key] 

del self.last_state[key] 

elif state == celery_states.FAILURE: 

self.fail(key) 

del self.tasks[key] 

del self.last_state[key] 

elif state == celery_states.REVOKED: 

self.fail(key) 

del self.tasks[key] 

del self.last_state[key] 

else: 

self.logger.info("Unexpected state: " + async.state) 

self.last_state[key] = async.state 

def end(self, synchronous=False): 

if synchronous: 

while any([ 

async.state not in celery_states.READY_STATES 

for async in self.tasks.values()]): 

time.sleep(5) 

self.sync()