breaking up base_executor into multiple files

airflow/bin/airflow

@@ -68,7 +68,7 @@ def run(args):
         dag = dag_pickle.get_object()
         task = dag.get_task(task_id=args.task_id)
 
-    # TODO: shouldn't the task instance run with the appropriate executor?
+    # TODO: add run_local and fire it with the right executor from run
     ti = TaskInstance(task, args.execution_date)
 
     # This is enough to fail the task instance

airflow/executors/__init__.py

@@ -1,6 +1,7 @@
-from base_executor import LocalExecutor
-from base_executor import SequentialExecutor
-from base_executor import CeleryExecutor
+from celery_executor import CeleryExecutor
+from local_executor import LocalExecutor
+from sequential_executor import SequentialExecutor
 
-DEFAULT_EXECUTOR = SequentialExecutor()
 # DEFAULT_EXECUTOR = CeleryExecutor()
+# DEFAULT_EXECUTOR = LocalExecutor()
+DEFAULT_EXECUTOR = SequentialExecutor()

airflow/executors/base_executor.py

@@ -1,10 +1,4 @@
 import logging
-import time
-from celery_worker import execute_command
-
-from airflow import settings
-from airflow.utils import State
-
 
 class BaseExecutor(object):
 
@@ -52,158 +46,3 @@ class BaseExecutor(object):
         """
         raise NotImplementedError()
 
-import multiprocessing
-import subprocess
-
-
-class LocalWorker(multiprocessing.Process):
-
-    def __init__(self, task_queue, result_queue):
-        multiprocessing.Process.__init__(self)
-        self.task_queue = task_queue
-        self.result_queue = result_queue
-
-    def run(self):
-        proc_name = self.name
-        while True:
-            key, command = self.task_queue.get()
-            if key is None:
-                # Received poison pill, no more tasks to run
-                self.task_queue.task_done()
-                break
-            BASE_FOLDER = settings.BASE_FOLDER
-            print command
-            command = (
-                "exec bash -c '"
-                "cd $AIRFLOW_HOME;\n" +
-                "source init.sh;\n" +
-                command +
-                "'"
-            ).format(**locals())
-            try:
-                sp = subprocess.Popen(command, shell=True).wait()
-            except Exception as e:
-                self.result_queue.put((key, State.FAILED))
-                raise e
-            self.result_queue.put((key, State.SUCCESS))
-            self.task_queue.task_done()
-            time.sleep(1)
-
-
-class LocalExecutor(BaseExecutor):
-
-    def __init__(self, parallelism=8):
-        super(LocalExecutor, self).__init__()
-        self.parallelism = parallelism
-
-    def start(self):
-        self.queue = multiprocessing.JoinableQueue()
-        self.result_queue = multiprocessing.Queue()
-        self.workers = [
-            LocalWorker(self.queue, self.result_queue)
-            for i in xrange(self.parallelism)
-        ]
-
-        for w in self.workers:
-            w.start()
-
-    def execute_async(self, key, command):
-        self.queue.put((key, command))
-
-    def heartbeat(self):
-        while not self.result_queue.empty():
-            results = self.result_queue.get()
-            self.change_state(*results)
-
-    def end(self):
-        # Sending poison pill to all worker
-        [self.queue.put((None, None)) for w in self.workers]
-        # Wait for commands to finish
-        self.queue.join()
-
-class CelerySubmitter(multiprocessing.Process):
-
-    def __init__(self, task_queue, result_queue):
-        multiprocessing.Process.__init__(self)
-        self.task_queue = task_queue
-        self.result_queue = result_queue
-
-    def run(self):
-        while True:
-            key, command = self.task_queue.get()
-            if command == None:
-                # Received poison pill, no more tasks to run
-                self.task_queue.task_done()
-                break
-            BASE_FOLDER = settings.BASE_FOLDER
-            command = (
-                "exec bash -c '"
-                "cd $AIRFLOW_HOME;\n" +
-                "source init.sh;\n" +
-                command +
-                "'"
-            ).format(**locals())
-
-            try:
-                res = execute_command.delay(command)
-                result = res.get()
-            except Exception as e:
-                self.result_queue.put((key, State.FAILED))
-                raise e
-            self.result_queue.put((key, State.SUCCESS))
-            self.task_queue.task_done()
-            time.sleep(1)
-
-
-class CeleryExecutor(BaseExecutor):
-    """ Submits the task to RabbitMQ, which is picked up and executed by a bunch
-        of worker processes """
-    def __init__(self, parallelism=1):
-        super(CeleryExecutor, self).__init__()
-        self.parallelism = parallelism
-
-    def start(self):
-        self.queue = multiprocessing.JoinableQueue()
-        self.result_queue = multiprocessing.Queue()
-        self.workers = [ CelerySubmitter(self.queue, self.result_queue) for i in xrange(self.parallelism) ]
-
-        for w in self.workers:
-            w.start()
-
-    def execute_async(self, key, command):
-        self.queue.put((key, command))
-
-    def heartbeat(self):
-        while not self.result_queue.empty():
-            results = self.result_queue.get()
-            self.change_state(*results)
-
-    def end(self):
-        # Sending poison pill to all worker
-        [self.queue.put(None) for w in self.workers]
-        self.queue.join()
-
-
-class SequentialExecutor(BaseExecutor):
-    """
-    Will only run one task instance at a time, can be used for debugging.
-    """
-    def __init__(self):
-        super(SequentialExecutor, self).__init__()
-        self.commands_to_run = []
-
-    def queue_command(self, key, command):
-        self.commands_to_run.append((key, command,))
-
-    def heartbeat(self):
-        for key, command in self.commands_to_run:
-            try:
-                sp = subprocess.Popen(command, shell=True).wait()
-            except Exception as e:
-                self.change_state(key, State.FAILED)
-                raise e
-            self.change_state(key, State.SUCCESS)
-        self.commands_to_run = []
-
-    def end(self):
-        pass

airflow/executors/celery_executor.py

@@ -0,0 +1,72 @@
+import multiprocessing
+import subprocess
+import time
+
+from airflow.executors.base_executor import BaseExecutor
+from airflow import settings
+from airflow.utils import State
+
+from celery_worker import execute_command
+
+class CeleryExecutor(BaseExecutor):
+    """ Submits the task to RabbitMQ, which is picked up and executed by a bunch
+        of worker processes """
+    def __init__(self, parallelism=1):
+        super(CeleryExecutor, self).__init__()
+        self.parallelism = parallelism
+
+    def start(self):
+        self.queue = multiprocessing.JoinableQueue()
+        self.result_queue = multiprocessing.Queue()
+        self.workers = [ CelerySubmitter(self.queue, self.result_queue) for i in xrange(self.parallelism) ]
+
+        for w in self.workers:
+            w.start()
+
+    def execute_async(self, key, command):
+        self.queue.put((key, command))
+
+    def heartbeat(self):
+        while not self.result_queue.empty():
+            results = self.result_queue.get()
+            self.change_state(*results)
+
+    def end(self):
+        # Sending poison pill to all worker
+        [self.queue.put(None) for w in self.workers]
+        self.queue.join()
+
+
+class CelerySubmitter(multiprocessing.Process):
+
+    def __init__(self, task_queue, result_queue):
+        multiprocessing.Process.__init__(self)
+        self.task_queue = task_queue
+        self.result_queue = result_queue
+
+    def run(self):
+        while True:
+            key, command = self.task_queue.get()
+            if command == None:
+                # Received poison pill, no more tasks to run
+                self.task_queue.task_done()
+                break
+            BASE_FOLDER = settings.BASE_FOLDER
+            command = (
+                "exec bash -c '"
+                "cd $AIRFLOW_HOME;\n" +
+                "source init.sh;\n" +
+                command +
+                "'"
+            ).format(**locals())
+
+            try:
+                res = execute_command.delay(command)
+                result = res.get()
+            except Exception as e:
+                self.result_queue.put((key, State.FAILED))
+                raise e
+            self.result_queue.put((key, State.SUCCESS))
+            self.task_queue.task_done()
+            time.sleep(1)
+

airflow/executors/local_executor.py

@@ -0,0 +1,74 @@
+import logging
+import multiprocessing
+import subprocess
+import time
+
+from airflow import settings
+from airflow.utils import State
+
+from airflow.executors.base_executor import BaseExecutor
+
+class LocalWorker(multiprocessing.Process):
+
+    def __init__(self, task_queue, result_queue):
+        multiprocessing.Process.__init__(self)
+        self.task_queue = task_queue
+        self.result_queue = result_queue
+
+    def run(self):
+        proc_name = self.name
+        while True:
+            key, command = self.task_queue.get()
+            if key is None:
+                # Received poison pill, no more tasks to run
+                self.task_queue.task_done()
+                break
+            BASE_FOLDER = settings.BASE_FOLDER
+            print command
+            command = (
+                "exec bash -c '"
+                "cd $AIRFLOW_HOME;\n" +
+                "source init.sh;\n" +
+                command +
+                "'"
+            ).format(**locals())
+            try:
+                sp = subprocess.Popen(command, shell=True).wait()
+            except Exception as e:
+                self.result_queue.put((key, State.FAILED))
+                raise e
+            self.result_queue.put((key, State.SUCCESS))
+            self.task_queue.task_done()
+            time.sleep(1)
+
+
+class LocalExecutor(BaseExecutor):
+
+    def __init__(self, parallelism=8):
+        super(LocalExecutor, self).__init__()
+        self.parallelism = parallelism
+
+    def start(self):
+        self.queue = multiprocessing.JoinableQueue()
+        self.result_queue = multiprocessing.Queue()
+        self.workers = [
+            LocalWorker(self.queue, self.result_queue)
+            for i in xrange(self.parallelism)
+        ]
+
+        for w in self.workers:
+            w.start()
+
+    def execute_async(self, key, command):
+        self.queue.put((key, command))
+
+    def heartbeat(self):
+        while not self.result_queue.empty():
+            results = self.result_queue.get()
+            self.change_state(*results)
+
+    def end(self):
+        # Sending poison pill to all worker
+        [self.queue.put((None, None)) for w in self.workers]
+        # Wait for commands to finish
+        self.queue.join()

airflow/executors/sequential_executor.py

@@ -0,0 +1,29 @@
+import subprocess
+
+from airflow.utils import State
+
+from airflow.executors.base_executor import BaseExecutor
+
+class SequentialExecutor(BaseExecutor):
+    """
+    Will only run one task instance at a time, can be used for debugging.
+    """
+    def __init__(self):
+        super(SequentialExecutor, self).__init__()
+        self.commands_to_run = []
+
+    def queue_command(self, key, command):
+        self.commands_to_run.append((key, command,))
+
+    def heartbeat(self):
+        for key, command in self.commands_to_run:
+            try:
+                sp = subprocess.Popen(command, shell=True).wait()
+            except Exception as e:
+                self.change_state(key, State.FAILED)
+                raise e
+            self.change_state(key, State.SUCCESS)
+        self.commands_to_run = []
+
+    def end(self):
+        pass

requirements.txt

@@ -1,3 +1,4 @@
+celery
 chartkick
 flake8
 flask
@@ -15,5 +16,3 @@ sphinx
 sphinx_rtd_theme
 sqlalchemy
 thrift
-celery
-hive-thrift-py