You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by da...@apache.org on 2016/07/31 19:50:09 UTC
[2/2] incubator-airflow git commit: [AIRFLOW-160] Parse DAG files
through child processes
[AIRFLOW-160] Parse DAG files through child processes
Instead of parsing the DAG definition files in the same process as the
scheduler, this change parses the files in a child process. This helps
to isolate the scheduler from bad user code.
Closes #1636 from plypaul/plypaul_schedule_by_file_rebase_master
Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/fdb7e949
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/fdb7e949
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/fdb7e949
Branch: refs/heads/master
Commit: fdb7e949140b735b8554ae5b22ad752e86f6ebaf
Parents: 835bcb6
Author: Paul Yang <pa...@airbnb.com>
Authored: Sun Jul 31 12:49:30 2016 -0700
Committer: Dan Davydov <da...@airbnb.com>
Committed: Sun Jul 31 12:49:39 2016 -0700
----------------------------------------------------------------------
airflow/bin/cli.py | 33 +-
airflow/configuration.py | 5 +
airflow/dag/__init__.py | 14 +
airflow/dag/base_dag.py | 96 ++
airflow/executors/base_executor.py | 2 +-
airflow/jobs.py | 1224 ++++++++++++++----
airflow/models.py | 273 +++-
airflow/settings.py | 36 +-
airflow/utils/dag_processing.py | 612 +++++++++
airflow/utils/db.py | 9 +-
scripts/ci/check-license.sh | 2 +-
scripts/ci/setup_env.sh | 12 +
setup.py | 2 +
tests/core.py | 33 +-
tests/dags/no_dags.py | 14 +
tests/dags_with_system_exit/a_system_exit.py | 29 +
.../b_test_scheduler_dags.py | 29 +
tests/dags_with_system_exit/c_system_exit.py | 29 +
tests/jobs.py | 192 +--
19 files changed, 2239 insertions(+), 407 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/bin/cli.py
----------------------------------------------------------------------
diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py
index 2d02f1e..0d35661 100755
--- a/airflow/bin/cli.py
+++ b/airflow/bin/cli.py
@@ -31,6 +31,8 @@ import daemon
from daemon.pidfile import TimeoutPIDLockFile
import signal
import sys
+import threading
+import traceback
import airflow
from airflow import jobs, settings
@@ -45,10 +47,28 @@ from airflow.exceptions import AirflowException
DAGS_FOLDER = os.path.expanduser(conf.get('core', 'DAGS_FOLDER'))
-def sigint_handler(signal, frame):
+def sigint_handler(sig, frame):
sys.exit(0)
+def sigquit_handler(sig, frame):
+ """Helps debug deadlocks by printing stacktraces when this gets a SIGQUIT
+ e.g. kill -s QUIT <PID> or CTRL+\
+ """
+ print("Dumping stack traces for all threads in PID {}".format(os.getpid()))
+ id_to_name = dict([(th.ident, th.name) for th in threading.enumerate()])
+ code = []
+ for thread_id, stack in sys._current_frames().items():
+ code.append("\n# Thread: {}({})"
+ .format(id_to_name.get(thread_id, ""), thread_id))
+ for filename, line_number, name, line in traceback.extract_stack(stack):
+ code.append('File: "{}", line {}, in {}'
+ .format((filename, line_number, name)))
+ if line:
+ code.append(" {}".format(line.strip()))
+ print("\n".join(code))
+
+
def setup_logging(filename):
root = logging.getLogger()
handler = logging.FileHandler(filename)
@@ -538,6 +558,7 @@ def scheduler(args):
job = jobs.SchedulerJob(
dag_id=args.dag_id,
subdir=process_subdir(args.subdir),
+ run_duration=args.run_duration,
num_runs=args.num_runs,
do_pickle=args.do_pickle)
@@ -561,6 +582,7 @@ def scheduler(args):
else:
signal.signal(signal.SIGINT, sigint_handler)
signal.signal(signal.SIGTERM, sigint_handler)
+ signal.signal(signal.SIGQUIT, sigquit_handler)
job.run()
@@ -918,6 +940,10 @@ class CLIFactory(object):
default=False),
# scheduler
'dag_id_opt': Arg(("-d", "--dag_id"), help="The id of the dag to run"),
+ 'run_duration': Arg(
+ ("-r", "--run-duration"),
+ default=None, type=int,
+ help="Set number of seconds to execute before exiting"),
'num_runs': Arg(
("-n", "--num_runs"),
default=None, type=int,
@@ -1056,8 +1082,9 @@ class CLIFactory(object):
}, {
'func': scheduler,
'help': "Start a scheduler instance",
- 'args': ('dag_id_opt', 'subdir', 'num_runs', 'do_pickle',
- 'pid', 'daemon', 'stdout', 'stderr', 'log_file'),
+ 'args': ('dag_id_opt', 'subdir', 'run_duration', 'num_runs',
+ 'do_pickle', 'pid', 'daemon', 'stdout', 'stderr',
+ 'log_file'),
}, {
'func': worker,
'help': "Start a Celery worker node",
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/configuration.py
----------------------------------------------------------------------
diff --git a/airflow/configuration.py b/airflow/configuration.py
index 5a380ae..65b482c 100644
--- a/airflow/configuration.py
+++ b/airflow/configuration.py
@@ -141,6 +141,11 @@ defaults = {
'scheduler_heartbeat_sec': 60,
'authenticate': False,
'max_threads': 2,
+ 'run_duration': 30 * 60,
+ 'dag_dir_list_interval': 5 * 60,
+ 'print_stats_interval': 30,
+ 'min_file_process_interval': 180,
+ 'child_process_log_directory': '/tmp/airflow/scheduler/logs'
},
'celery': {
'broker_url': 'sqla+mysql://airflow:airflow@localhost:3306/airflow',
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/dag/__init__.py
----------------------------------------------------------------------
diff --git a/airflow/dag/__init__.py b/airflow/dag/__init__.py
new file mode 100644
index 0000000..a84b6da
--- /dev/null
+++ b/airflow/dag/__init__.py
@@ -0,0 +1,14 @@
+# -*- 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.
+#
\ No newline at end of file
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/dag/base_dag.py
----------------------------------------------------------------------
diff --git a/airflow/dag/base_dag.py b/airflow/dag/base_dag.py
new file mode 100644
index 0000000..83ecfb9
--- /dev/null
+++ b/airflow/dag/base_dag.py
@@ -0,0 +1,96 @@
+# -*- 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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from abc import ABCMeta, abstractmethod, abstractproperty
+
+
+class BaseDag(object):
+ """
+ Base DAG object that both the SimpleDag and DAG inherit.
+ """
+ __metaclass__ = ABCMeta
+
+ @abstractproperty
+ def dag_id(self):
+ """
+ :return: the DAG ID
+ :rtype: unicode
+ """
+ raise NotImplementedError()
+
+ @abstractproperty
+ def task_ids(self):
+ """
+ :return: A list of task IDs that are in this DAG
+ :rtype: List[unicode]
+ """
+ raise NotImplementedError()
+
+ @abstractproperty
+ def full_filepath(self):
+ """
+ :return: The absolute path to the file that contains this DAG's definition
+ :rtype: unicode
+ """
+ raise NotImplementedError()
+
+ @abstractmethod
+ def concurrency(self):
+ """
+ :return: maximum number of tasks that can run simultaneously from this DAG
+ :rtype: int
+ """
+ raise NotImplementedError()
+
+ @abstractmethod
+ def is_paused(self):
+ """
+ :return: whether this DAG is paused or not
+ :rtype: bool
+ """
+ raise NotImplementedError()
+
+ @abstractmethod
+ def pickle_id(self):
+ """
+ :return: The pickle ID for this DAG, if it has one. Otherwise None.
+ :rtype: unicode
+ """
+ raise NotImplementedError
+
+
+class BaseDagBag(object):
+ """
+ Base object that both the SimpleDagBag and DagBag inherit.
+ """
+ @abstractproperty
+ def dag_ids(self):
+ """
+ :return: a list of DAG IDs in this bag
+ :rtype: List[unicode]
+ """
+ raise NotImplementedError()
+
+ @abstractmethod
+ def get_dag(self, dag_id):
+ """
+ :return: whether the task exists in this bag
+ :rtype: BaseDag
+ """
+ raise NotImplementedError()
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/executors/base_executor.py
----------------------------------------------------------------------
diff --git a/airflow/executors/base_executor.py b/airflow/executors/base_executor.py
index ca63443..d6a06d8 100644
--- a/airflow/executors/base_executor.py
+++ b/airflow/executors/base_executor.py
@@ -102,7 +102,7 @@ class BaseExecutor(LoggingMixin):
# TODO(jlowin) without a way to know what Job ran which tasks,
# there is a danger that another Job started running a task
# that was also queued to this executor. This is the last chance
- # to check if that hapened. The most probable way is that a
+ # to check if that happened. The most probable way is that a
# Scheduler tried to run a task that was originally queued by a
# Backfill. This fix reduces the probability of a collision but
# does NOT eliminate it.
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/jobs.py
----------------------------------------------------------------------
diff --git a/airflow/jobs.py b/airflow/jobs.py
index 77f34ee..b391415 100644
--- a/airflow/jobs.py
+++ b/airflow/jobs.py
@@ -19,30 +19,45 @@ from __future__ import unicode_literals
from past.builtins import basestring
from collections import defaultdict, Counter
-from datetime import datetime, timedelta
+
+from datetime import datetime
+
import getpass
import logging
import socket
import subprocess
import multiprocessing
-import math
+import os
+import signal
+import sys
+import threading
+import time
from time import sleep
+import psutil
from sqlalchemy import Column, Integer, String, DateTime, func, Index, or_
from sqlalchemy.orm.session import make_transient
+from tabulate import tabulate
from airflow import executors, models, settings
from airflow import configuration as conf
from airflow.exceptions import AirflowException
+from airflow.models import DagRun
+from airflow.settings import Stats
from airflow.utils.state import State
from airflow.utils.db import provide_session, pessimistic_connection_handling
+from airflow.utils.dag_processing import (AbstractDagFileProcessor,
+ DagFileProcessorManager,
+ SimpleDag,
+ SimpleDagBag,
+ list_py_file_paths)
from airflow.utils.email import send_email
from airflow.utils.logging import LoggingMixin
from airflow.utils import asciiart
-from airflow.settings import Stats
-DagRun = models.DagRun
+
Base = models.Base
+DagRun = models.DagRun
ID_LEN = models.ID_LEN
Stats = settings.Stats
@@ -186,27 +201,247 @@ class BaseJob(Base, LoggingMixin):
raise NotImplementedError("This method needs to be overridden")
+class DagFileProcessor(AbstractDagFileProcessor):
+ """Helps call SchedulerJob.process_file() in a separate process."""
+
+ # Counter that increments everytime an instance of this class is created
+ class_creation_counter = 0
+
+ def __init__(self, file_path, pickle_dags, dag_id_white_list, log_file):
+ """
+ :param file_path: a Python file containing Airflow DAG definitions
+ :type file_path: unicode
+ :param pickle_dags: whether to serialize the DAG objects to the DB
+ :type pickle_dags: bool
+ :param dag_id_whitelist: If specified, only look at these DAG ID's
+ :type dag_id_whitelist: list[unicode]
+ :param log_file: the path to the file where log lines should be output
+ :type log_file: unicode
+ """
+ self._file_path = file_path
+ self._log_file = log_file
+ # Queue that's used to pass results from the child process.
+ self._result_queue = multiprocessing.Queue()
+ # The process that was launched to process the given .
+ self._process = None
+ self._dag_id_white_list = dag_id_white_list
+ self._pickle_dags = pickle_dags
+ # The result of Scheduler.process_file(file_path).
+ self._result = None
+ # Whether the process is done running.
+ self._done = False
+ # When the process started.
+ self._start_time = None
+ # This ID is use to uniquely name the process / thread that's launched
+ # by this processor instance
+ self._instance_id = DagFileProcessor.class_creation_counter
+ DagFileProcessor.class_creation_counter += 1
+
+ @property
+ def file_path(self):
+ return self._file_path
+
+ @property
+ def log_file(self):
+ return self._log_file
+
+ @staticmethod
+ def _launch_process(result_queue,
+ file_path,
+ pickle_dags,
+ dag_id_white_list,
+ thread_name,
+ log_file):
+ """
+ Launch a process to process the given file.
+
+ :param result_queue: the queue to use for passing back the result
+ :type result_queue: multiprocessing.Queue
+ :param file_path: the file to process
+ :type file_path: unicode
+ :param pickle_dags: whether to pickle the DAGs found in the file and
+ save them to the DB
+ :type pickle_dags: bool
+ :param dag_id_white_list: if specified, only examine DAG ID's that are
+ in this list
+ :type dag_id_white_list: list[unicode]
+ :param thread_name: the name to use for the process that is launched
+ :type thread_name: unicode
+ :param log_file: the logging output for the process should be directed
+ to this file
+ :type log_file: unicode
+ :return: the process that was launched
+ :rtype: multiprocessing.Process
+ """
+ def helper():
+ # This helper runs in the newly created process
+
+ # Re-direct stdout and stderr to a separate log file. Otherwise,
+ # the main log becomes too hard to read. No buffering to enable
+ # responsive file tailing
+ parent_dir, _ = os.path.split(log_file)
+
+ # Create the parent directory for the log file if necessary.
+ if not os.path.isdir(parent_dir):
+ os.makedirs(parent_dir)
+
+ f = open(log_file, "a")
+ original_stdout = sys.stdout
+ original_stderr = sys.stderr
+
+ sys.stdout = f
+ sys.stderr = f
+
+ try:
+ # Re-configure logging to use the new output streams
+ log_format = settings.LOG_FORMAT_WITH_THREAD_NAME
+ settings.configure_logging(log_format=log_format)
+ # Re-configure the ORM engine as there are issues with multiple processes
+ settings.configure_orm()
+
+ # Change the thread name to differentiate log lines. This is
+ # really a separate process, but changing the name of the
+ # process doesn't work, so changing the thread name instead.
+ threading.current_thread().name = thread_name
+ start_time = time.time()
+
+ logging.info("Started process (PID=%s) to work on %s",
+ os.getpid(),
+ file_path)
+ scheduler_job = SchedulerJob(dag_ids=dag_id_white_list)
+ result = scheduler_job.process_file(file_path,
+ pickle_dags)
+ result_queue.put(result)
+ end_time = time.time()
+ logging.info("Processing %s took %.3f seconds",
+ file_path,
+ end_time - start_time)
+ except:
+ # Log exceptions through the logging framework.
+ logging.exception("Got an exception! Propagating...")
+ raise
+ finally:
+ sys.stdout = original_stdout
+ sys.stderr = original_stderr
+ f.close()
+
+ p = multiprocessing.Process(target=helper,
+ args=(),
+ name="{}-Process".format(thread_name))
+ p.start()
+ return p
+
+ def start(self):
+ """
+ Launch the process and start processing the DAG.
+ """
+ self._process = DagFileProcessor._launch_process(
+ self._result_queue,
+ self.file_path,
+ self._pickle_dags,
+ self._dag_id_white_list,
+ "DagFileProcessor{}".format(self._instance_id),
+ self.log_file)
+ self._start_time = datetime.now()
+
+ def terminate(self, sigkill=False):
+ """
+ Terminate (and then kill) the process launched to process the file.
+ :param sigkill: whether to issue a SIGKILL if SIGTERM doesn't work.
+ :type sigkill: bool
+ """
+ if self._process is None:
+ raise AirflowException("Tried to call stop before starting!")
+ # The queue will likely get corrupted, so remove the reference
+ self._result_queue = None
+ self._process.terminate()
+ # Arbitrarily wait 5s for the process to die
+ self._process.join(5)
+ if sigkill and self._process.is_alive():
+ logging.warn("Killing PID %s", self._process.pid)
+ os.kill(self._process.pid, signal.SIGKILL)
+
+ @property
+ def pid(self):
+ """
+ :return: the PID of the process launched to process the given file
+ :rtype: int
+ """
+ if self._process is None:
+ raise AirflowException("Tried to get PID before starting!")
+ return self._process.pid
+
+ @property
+ def exit_code(self):
+ """
+ After the process is finished, this can be called to get the return code
+ :return: the exit code of the process
+ :rtype: int
+ """
+ if not self._done:
+ raise AirflowException("Tried to call retcode before process was finished!")
+ return self._process.exitcode
+
+ @property
+ def done(self):
+ """
+ Check if the process launched to process this file is done.
+ :return: whether the process is finished running
+ :rtype: bool
+ """
+ if self._process is None:
+ raise AirflowException("Tried to see if it's done before starting!")
+
+ if self._done:
+ return True
+
+ if not self._result_queue.empty():
+ self._result = self._result_queue.get_nowait()
+ self._done = True
+ logging.debug("Waiting for %s", self._process)
+ self._process.join()
+ return True
+
+ # Potential error case when process dies
+ if not self._process.is_alive():
+ self._done = True
+ # Get the object from the queue or else join() can hang.
+ if not self._result_queue.empty():
+ self._result = self._result_queue.get_nowait()
+ logging.debug("Waiting for %s", self._process)
+ self._process.join()
+ return True
+
+ return False
+
+ @property
+ def result(self):
+ """
+ :return: result of running SchedulerJob.process_file()
+ :rtype: SimpleDag
+ """
+ if not self.done:
+ raise AirflowException("Tried to get the result before it's done!")
+ return self._result
+
+ @property
+ def start_time(self):
+ """
+ :return: when this started to process the file
+ :rtype: datetime
+ """
+ if self._start_time is None:
+ raise AirflowException("Tried to get start time before it started!")
+ return self._start_time
+
+
class SchedulerJob(BaseJob):
"""
- This SchedulerJob runs indefinitely and constantly schedules the jobs
+ This SchedulerJob runs for a specific time interval and schedules the jobs
that are ready to run. It figures out the latest runs for each
- task and see if the dependencies for the next schedules are met.
- If so it triggers the task instance. It does this for each task
- in each DAG and repeats.
-
- :param dag_id: to run the scheduler for a single specific DAG
- :type dag_id: string
- :param subdir: to search for DAG under a certain folder only
- :type subdir: string
- :param test_mode: used for unit testing this class only, runs a single
- schedule run
- :type test_mode: bool
- :param refresh_dags_every: force refresh the DAG definition every N
- runs, as specified here
- :type refresh_dags_every: int
- :param do_pickle: to pickle the DAG object and send over to workers
- for non-local executors
- :type do_pickle: bool
+ task and sees if the dependencies for the next schedules are met.
+ If so, it creates appropriate TaskInstances and sends run commands to the
+ executor. It does this for each task in each DAG and repeats.
"""
__mapper_args__ = {
@@ -217,13 +452,32 @@ class SchedulerJob(BaseJob):
self,
dag_id=None,
dag_ids=None,
- subdir=None,
- test_mode=False,
- refresh_dags_every=10,
- num_runs=None,
+ subdir=models.DAGS_FOLDER,
+ num_runs=-1,
+ file_process_interval=conf.getint('scheduler',
+ 'min_file_process_interval'),
+ processor_poll_interval=1.0,
+ run_duration=None,
do_pickle=False,
*args, **kwargs):
-
+ """
+ :param dag_id: if specified, only schedule tasks with this DAG ID
+ :type dag_id: unicode
+ :param dag_ids: if specified, only schedule tasks with these DAG IDs
+ :type dag_ids: list[unicode]
+ :param subdir: directory containing Python files with Airflow DAG
+ definitions, or a specific path to a file
+ :type subdir: unicode
+ :param num_runs: The number of times to try to schedule each DAG file.
+ -1 for unlimited within the run_duration.
+ :param processor_poll_interval: The number of seconds to wait between
+ polls of running processors
+ :param run_duration: how long to run (in seconds) before exiting
+ :type run_duration: int
+ :param do_pickle: once a DAG object is obtained by executing the Python
+ file, whether to serialize the DAG object to the DB
+ :type do_pickle: bool
+ """
# for BaseJob compatibility
self.dag_id = dag_id
self.dag_ids = [dag_id] if dag_id else []
@@ -232,21 +486,38 @@ class SchedulerJob(BaseJob):
self.subdir = subdir
- if test_mode:
- self.num_runs = 1
- else:
- self.num_runs = num_runs
+ self.num_runs = num_runs
+ self.run_duration = run_duration
+ self._processor_poll_interval = processor_poll_interval
- self.refresh_dags_every = refresh_dags_every
self.do_pickle = do_pickle
super(SchedulerJob, self).__init__(*args, **kwargs)
self.heartrate = conf.getint('scheduler', 'SCHEDULER_HEARTBEAT_SEC')
self.max_threads = min(conf.getint('scheduler', 'max_threads'), multiprocessing.cpu_count())
+ self.using_sqlite = False
if 'sqlite' in conf.get('core', 'sql_alchemy_conn'):
if self.max_threads > 1:
self.logger.error("Cannot use more than 1 thread when using sqlite. Setting max_threads to 1")
self.max_threads = 1
+ self.using_sqlite = True
+
+ # How often to scan the DAGs directory for new files. Default to 5 minutes.
+ self.dag_dir_list_interval = conf.getint('scheduler',
+ 'dag_dir_list_interval')
+ # How often to print out DAG file processing stats to the log. Default to
+ # 30 seconds.
+ self.print_stats_interval = conf.getint('scheduler',
+ 'print_stats_interval')
+ # Parse and schedule each file no faster than this interval. Default
+ # to 3 minutes.
+ self.file_process_interval = file_process_interval
+ # Directory where log files for the processes that scheduled the DAGs reside
+ self.child_process_log_directory = conf.get('scheduler',
+ 'child_process_log_directory')
+ if run_duration is None:
+ self.run_duration = conf.getint('scheduler',
+ 'run_duration')
@provide_session
def manage_slas(self, dag, session=None):
@@ -296,7 +567,8 @@ class SchedulerJob(BaseJob):
slas = (
session
.query(SlaMiss)
- .filter(SlaMiss.email_sent.is_(False) or SlaMiss.notification_sent.is_(False))
+ .filter(or_(SlaMiss.email_sent == False,
+ SlaMiss.notification_sent == False))
.filter(SlaMiss.dag_id == dag.dag_id)
.all()
)
@@ -370,23 +642,45 @@ class SchedulerJob(BaseJob):
session.commit()
session.close()
- def import_errors(self, dagbag):
- session = settings.Session()
- session.query(models.ImportError).delete()
+ @staticmethod
+ def record_import_errors(session, dagbag):
+ """
+ For the DAGs in the given DagBag, record any associated import errors.
+ These are usually displayed through the Airflow UI so that users know
+ that there are issues parsing DAGs.
+
+ :param session: session for ORM operations
+ :type session: sqlalchemy.orm.session.Session
+ :param dagbag: DagBag containing DAGs with import errors
+ :type dagbag: models.Dagbag
+ """
for filename, stacktrace in list(dagbag.import_errors.items()):
+ session.query(models.ImportError).filter(
+ models.ImportError.filename == filename
+ ).delete()
session.add(models.ImportError(
filename=filename, stacktrace=stacktrace))
session.commit()
+ @staticmethod
+ def clear_import_errors(session):
+ """
+ Remove all the known import errors from the DB.
+
+ :param session: session for ORM operations
+ :type session: sqlalchemy.orm.session.Session
+ """
+ session.query(models.ImportError).delete()
+ session.commit()
+
@provide_session
- def schedule_dag(self, dag, session=None):
+ def create_dag_run(self, dag, session=None):
"""
This method checks whether a new DagRun needs to be created
for a DAG based on scheduling interval
Returns DagRun if one is scheduled. Otherwise returns None.
"""
if dag.schedule_interval:
- DagRun = models.DagRun
active_runs = DagRun.find(
dag_id=dag.dag_id,
state=State.RUNNING,
@@ -475,52 +769,24 @@ class SchedulerJob(BaseJob):
)
return next_run
- def process_dag(self, dag, queue):
+ def _process_task_instances(self, dag, queue):
"""
- This method schedules a single DAG by looking at the latest
- run for each task and attempting to schedule the following run.
-
- As multiple schedulers may be running for redundancy, this
- function takes a lock on the DAG and timestamps the last run
- in ``last_scheduler_run``.
+ This method schedules the tasks for a single DAG by looking at the
+ active DAG runs and adding task instances that should run to the
+ queue.
"""
DagModel = models.DagModel
session = settings.Session()
- # picklin'
- pickle_id = None
- if self.do_pickle and self.executor.__class__ not in (
- executors.LocalExecutor, executors.SequentialExecutor):
- pickle_id = dag.pickle(session).id
-
- # obtain db lock
- db_dag = session.query(DagModel).filter_by(
- dag_id=dag.dag_id
- ).with_for_update().one()
-
- last_scheduler_run = db_dag.last_scheduler_run or datetime(2000, 1, 1)
- secs_since_last = (datetime.now() - last_scheduler_run).total_seconds()
-
- if secs_since_last < self.heartrate:
- # release db lock
- session.commit()
- session.close()
- return None
-
- # Release the db lock
- # the assumption here is that process_dag will take less
- # time than self.heartrate otherwise we might unlock too
- # quickly and this should moved below, but that would increase
- # the time the record is locked and is blocking for other calls.
- db_dag.last_scheduler_run = datetime.now()
- session.commit()
-
# update the state of the previously active dag runs
dag_runs = DagRun.find(dag_id=dag.dag_id, state=State.RUNNING, session=session)
active_dag_runs = []
for run in dag_runs:
+ self.logger.info("Examining DAG run {}".format(run))
# do not consider runs that are executed in the future
if run.execution_date > datetime.now():
+ self.logging.error("Execution date is in future: {}"
+ .format(run.execution_date))
continue
# todo: run.dag is transient but needs to be set
@@ -533,6 +799,7 @@ class SchedulerJob(BaseJob):
active_dag_runs.append(run)
for run in active_dag_runs:
+ self.logger.debug("Examining active DAG run {}".format(run))
# this needs a fresh session sometimes tis get detached
tis = run.get_task_instances(state=(State.NONE,
State.UP_FOR_RETRY))
@@ -552,41 +819,122 @@ class SchedulerJob(BaseJob):
if ti.is_runnable(flag_upstream_failed=True):
self.logger.debug('Queuing task: {}'.format(ti))
- queue.put((ti.key, pickle_id))
+ queue.append(ti.key)
session.close()
@provide_session
- def prioritize_queued(self, session, executor, dagbag):
- # Prioritizing queued task instances
+ def _change_state_for_tis_without_dagrun(self,
+ simple_dag_bag,
+ old_states,
+ new_state,
+ session=None):
+ """
+ For all DAG IDs in the SimpleDagBag, look for task instances in the
+ old_states and set them to new_state if the corresponding DagRun
+ exists but is not in the running state. This normally should not
+ happen, but it can if the state of DagRuns are changed manually.
+
+ :param old_states: examine TaskInstances in this state
+ :type old_state: list[State]
+ :param new_state: set TaskInstances to this state
+ :type new_state: State
+ :param simple_dag_bag: TaskInstances associated with DAGs in the
+ simple_dag_bag and with states in the old_state will be examined
+ :type simple_dag_bag: SimpleDagBag
+ """
- pools = {p.pool: p for p in session.query(models.Pool).all()}
+ task_instances_to_change = (
+ session
+ .query(models.TaskInstance)
+ .filter(models.TaskInstance.dag_id.in_(simple_dag_bag.dag_ids))
+ .filter(models.TaskInstance.state.in_(old_states))
+ .all()
+ )
+ """:type: list[TaskInstance]"""
+
+ for task_instance in task_instances_to_change:
+ dag_runs = DagRun.find(dag_id=task_instance.dag_id,
+ execution_date=task_instance.execution_date,
+ )
+
+ if len(dag_runs) == 0:
+ self.logger.warn("DagRun for %s %s does not exist",
+ task_instance.dag_id,
+ task_instance.execution_date)
+ continue
+
+ # There should only be one DAG run. Add some logging info if this
+ # is not the case for later debugging.
+ if len(dag_runs) > 1:
+ self.logger.warn("Multiple DagRuns found for {} {}: {}"
+ .format(task_instance.dag_id,
+ task_instance.execution_date,
+ dag_runs))
+
+ dag_is_running = True
+ for dag_run in dag_runs:
+ if dag_run.state == State.RUNNING:
+ dag_is_running = True
+ break
+
+ if not dag_is_running:
+ self.logger.warn("Setting {} to state={} as it does not have "
+ "a DagRun in the {} state"
+ .format(task_instance,
+ new_state,
+ State.RUNNING))
+ task_instance.state = new_state
+ session.merge(task_instance)
+ session.commit()
+
+ @provide_session
+ def _execute_task_instances(self,
+ simple_dag_bag,
+ states,
+ session=None):
+ """
+ Fetches task instances from ORM in the specified states, figures
+ out pool limits, and sends them to the executor for execution.
+
+ :param simple_dag_bag: TaskInstances associated with DAGs in the
+ simple_dag_bag will be fetched from the DB and executed
+ :type simple_dag_bag: SimpleDagBag
+ :param executor: the executor that runs task instances
+ :type executor: BaseExecutor
+ :param states: Execute TaskInstances in these states
+ :type states: Tuple[State]
+ :return: None
+ """
+ # Get all the relevant task instances
TI = models.TaskInstance
- queued_tis = (
- session.query(TI)
- .filter(TI.state == State.QUEUED)
+ task_instances_to_examine = (
+ session
+ .query(TI)
+ .filter(TI.dag_id.in_(simple_dag_bag.dag_ids))
+ .filter(TI.state.in_(states))
.all()
)
- self.logger.info(
- "Prioritizing {} queued jobs".format(len(queued_tis)))
- session.expunge_all()
- d = defaultdict(list)
- for ti in queued_tis:
- if ti.dag_id not in dagbag.dags:
- self.logger.info(
- "DAG no longer in dagbag, deleting {}".format(ti))
- session.delete(ti)
- session.commit()
- elif not dagbag.dags[ti.dag_id].has_task(ti.task_id):
- self.logger.info(
- "Task no longer exists, deleting {}".format(ti))
- session.delete(ti)
- session.commit()
- else:
- d[ti.pool].append(ti)
- dag_blacklist = set(dagbag.paused_dags())
- for pool, tis in list(d.items()):
+ # Put one task instance on each line
+ if len(task_instances_to_examine) == 0:
+ self.logger.info("No tasks to send to the executor")
+ return
+
+ task_instance_str = "\n\t".join(
+ ["{}".format(x) for x in task_instances_to_examine])
+ self.logger.info("Tasks up for execution:\n\t{}".format(task_instance_str))
+
+ # Get the pool settings
+ pools = {p.pool: p for p in session.query(models.Pool).all()}
+
+ pool_to_task_instances = defaultdict(list)
+ for task_instance in task_instances_to_examine:
+ pool_to_task_instances[task_instance.pool].append(task_instance)
+
+ # Go through each pool, and queue up a task for execution if there are
+ # any open slots in the pool.
+ for pool, task_instances in pool_to_task_instances.items():
if not pool:
# Arbitrary:
# If queued outside of a pool, trigger no more than
@@ -595,78 +943,218 @@ class SchedulerJob(BaseJob):
else:
open_slots = pools[pool].open_slots(session=session)
- queue_size = len(tis)
- self.logger.info("Pool {pool} has {open_slots} slots, {queue_size} "
+ num_queued = len(task_instances)
+ self.logger.info("Figuring out tasks to run in Pool(name={pool}) "
+ "with {open_slots} open slots and {num_queued} "
"task instances in queue".format(**locals()))
+
if open_slots <= 0:
continue
- tis = sorted(
- tis, key=lambda ti: (-ti.priority_weight, ti.start_date))
- for ti in tis:
- if open_slots <= 0:
- continue
- task = None
- try:
- task = dagbag.dags[ti.dag_id].get_task(ti.task_id)
- except:
- self.logger.error("Queued task {} seems gone".format(ti))
- session.delete(ti)
- session.commit()
- continue
- if not task:
- continue
+ priority_sorted_task_instances = sorted(
+ task_instances, key=lambda ti: (-ti.priority_weight, ti.execution_date))
- ti.task = task
+ # DAG IDs with running tasks that equal the concurrency limit of the dag
+ dag_id_to_running_task_count = {}
- # picklin'
- dag = dagbag.dags[ti.dag_id]
- pickle_id = None
- if self.do_pickle and self.executor.__class__ not in (
- executors.LocalExecutor,
- executors.SequentialExecutor):
- self.logger.info("Pickling DAG {}".format(dag))
- pickle_id = dag.pickle(session).id
+ for task_instance in priority_sorted_task_instances:
+ if open_slots <= 0:
+ self.logger.info("No more slots free")
+ # Can't schedule any more since there are no more open slots.
+ break
- if dag.dag_id in dag_blacklist:
- continue
- if dag.concurrency_reached:
- dag_blacklist.add(dag.dag_id)
+ if simple_dag_bag.get_dag(task_instance.dag_id).is_paused:
+ self.logger.info("Not executing queued {} since {} is paused"
+ .format(task_instance, task_instance.dag_id))
continue
- if ti.are_dependencies_met():
- executor.queue_task_instance(ti, pickle_id=pickle_id)
- open_slots -= 1
- else:
- session.delete(ti)
- session.commit()
+
+ # Check to make sure that the task concurrency of the DAG hasn't been
+ # reached.
+ dag_id = task_instance.dag_id
+
+ if dag_id not in dag_id_to_running_task_count:
+ dag_id_to_running_task_count[dag_id] = \
+ DagRun.get_running_tasks(
+ session,
+ dag_id,
+ simple_dag_bag.get_dag(dag_id).task_ids)
+
+ current_task_concurrency = dag_id_to_running_task_count[dag_id]
+ task_concurrency_limit = simple_dag_bag.get_dag(dag_id).concurrency
+ self.logger.info("DAG {} has {}/{} running tasks"
+ .format(dag_id,
+ current_task_concurrency,
+ task_concurrency_limit))
+ if current_task_concurrency > task_concurrency_limit:
+ self.logger.info("Not executing {} since the number "
+ "of tasks running from DAG {} is >= to the "
+ "DAG's task concurrency limit of {}"
+ .format(task_instance,
+ dag_id,
+ task_concurrency_limit))
continue
- ti.task = task
+ command = TI.generate_command(
+ task_instance.dag_id,
+ task_instance.task_id,
+ task_instance.execution_date,
+ local=True,
+ mark_success=False,
+ force=False,
+ ignore_dependencies=False,
+ ignore_depends_on_past=False,
+ pool=task_instance.pool,
+ file_path=simple_dag_bag.get_dag(task_instance.dag_id).full_filepath,
+ pickle_id=simple_dag_bag.get_dag(task_instance.dag_id).pickle_id)
+
+ priority = task_instance.priority_weight
+ queue = task_instance.queue
+ self.logger.info("Sending to executor {} with priority {} and queue {}"
+ .format(task_instance.key, priority, queue))
+
+ # Set the state to queued
+ self.logger.info("Setting state of {} to {}".format(
+ task_instance.key, State.QUEUED))
+ task_instance.state = State.QUEUED
+ task_instance.queued_dttm = (datetime.now()
+ if not task_instance.queued_dttm
+ else task_instance.queued_dttm)
+ session.merge(task_instance)
session.commit()
- def _split(self, items, size):
- """
- This function splits a list of items into chunks of int size.
- _split([1,2,3,4,5,6], 3) becomes [[1,2,3],[4,5,6]]
- """
- size = max(1, size)
- return [items[i:i + size] for i in range(0, len(items), size)]
+ # These attributes will be lost after the object expires, so save them.
+ task_id_ = task_instance.task_id
+ dag_id_ = task_instance.dag_id
+ execution_date_ = task_instance.execution_date
+ make_transient(task_instance)
+ task_instance.task_id = task_id_
+ task_instance.dag_id = dag_id_
+ task_instance.execution_date = execution_date_
+
+ self.executor.queue_command(
+ task_instance,
+ command,
+ priority=priority,
+ queue=queue)
- def _do_dags(self, dagbag, dags, tis_out):
+ open_slots -= 1
+
+ def _process_dags(self, dagbag, dags, tis_out):
"""
- Iterates over the dags and schedules and processes them
+ Iterates over the dags and processes them. Processing includes:
+
+ 1. Create appropriate DagRun(s) in the DB.
+ 2. Create appropriate TaskInstance(s) in the DB.
+ 3. Send emails for tasks that have missed SLAs.
+
+ :param dagbag: a collection of DAGs to process
+ :type dagbag: models.DagBag
+ :param dags: the DAGs from the DagBag to process
+ :type dags: DAG
+ :param tis_out: A queue to add generated TaskInstance objects
+ :type tis_out: multiprocessing.Queue[TaskInstance]
+ :return: None
"""
for dag in dags:
- self.logger.debug("Scheduling {}".format(dag.dag_id))
dag = dagbag.get_dag(dag.dag_id)
+ if dag.is_paused:
+ self.logger.info("Not processing DAG {} since it's paused"
+ .format(dag.dag_id))
+ continue
+
if not dag:
+ self.logger.error("DAG ID {} was not found in the DagBag")
continue
- try:
- self.schedule_dag(dag)
- self.process_dag(dag, tis_out)
- self.manage_slas(dag)
- except Exception as e:
- self.logger.exception(e)
+
+ self.logger.info("Processing {}".format(dag.dag_id))
+
+ dag_run = self.create_dag_run(dag)
+ if dag_run:
+ self.logger.info("Created {}".format(dag_run))
+ self._process_task_instances(dag, tis_out)
+ self.manage_slas(dag)
+
+ def _process_executor_events(self):
+ """
+ Respond to executor events.
+
+ :param executor: the executor that's running the task instances
+ :type executor: BaseExecutor
+ :return: None
+ """
+ for key, executor_state in list(self.executor.get_event_buffer().items()):
+ dag_id, task_id, execution_date = key
+ self.logger.info("Executor reports {}.{} execution_date={} as {}"
+ .format(dag_id,
+ task_id,
+ execution_date,
+ executor_state))
+
+ def _log_file_processing_stats(self,
+ known_file_paths,
+ processor_manager):
+ """
+ Print out stats about how files are getting processed.
+
+ :param known_file_paths: a list of file paths that may contain Airflow
+ DAG definitions
+ :type known_file_paths: list[unicode]
+ :param processor_manager: manager for the file processors
+ :type stats: DagFileProcessorManager
+ :return: None
+ """
+
+ # File Path: Path to the file containing the DAG definition
+ # PID: PID associated with the process that's processing the file. May
+ # be empty.
+ # Runtime: If the process is currently running, how long it's been
+ # running for in seconds.
+ # Last Runtime: If the process ran before, how long did it take to
+ # finish in seconds
+ # Last Run: When the file finished processing in the previous run.
+ headers = ["File Path",
+ "PID",
+ "Runtime",
+ "Last Runtime",
+ "Last Run"]
+
+ rows = []
+ for file_path in known_file_paths:
+ last_runtime = processor_manager.get_last_runtime(file_path)
+ processor_pid = processor_manager.get_pid(file_path)
+ processor_start_time = processor_manager.get_start_time(file_path)
+ runtime = ((datetime.now() - processor_start_time).total_seconds()
+ if processor_start_time else None)
+ last_run = processor_manager.get_last_finish_time(file_path)
+
+ rows.append((file_path,
+ processor_pid,
+ runtime,
+ last_runtime,
+ last_run))
+
+ # Sort by longest last runtime. (Can't sort None values in python3)
+ rows = sorted(rows, key=lambda x: x[3] or 0.0)
+
+ formatted_rows = []
+ for file_path, pid, runtime, last_runtime, last_run in rows:
+ formatted_rows.append((file_path,
+ pid,
+ "{:.2f}s".format(runtime)
+ if runtime else None,
+ "{:.2f}s".format(last_runtime)
+ if last_runtime else None,
+ last_run.strftime("%Y-%m-%dT%H:%M:%S")
+ if last_run else None))
+ log_str = ("\n" +
+ "=" * 80 +
+ "\n" +
+ "DAG File Processing Stats\n\n" +
+ tabulate(formatted_rows, headers=headers) +
+ "\n" +
+ "=" * 80)
+
+ self.logger.info(log_str)
@provide_session
def _reset_state_for_orphaned_tasks(self, dag_run, session=None):
@@ -680,27 +1168,110 @@ class SchedulerJob(BaseJob):
# also consider running as the state might not have changed in the db yet
running = self.executor.running
- tis = dag_run.get_task_instances(state=State.SCHEDULED, session=session)
+ tis = list()
+ tis.extend(dag_run.get_task_instances(state=State.SCHEDULED, session=session))
+ tis.extend(dag_run.get_task_instances(state=State.QUEUED, session=session))
+
for ti in tis:
if ti.key not in queued_tis and ti.key not in running:
- ti.state = State.NONE
self.logger.debug("Rescheduling orphaned task {}".format(ti))
-
+ ti.state = State.NONE
session.commit()
def _execute(self):
- session = settings.Session()
- TI = models.TaskInstance
-
+ self.logger.info("Starting the scheduler")
pessimistic_connection_handling()
logging.basicConfig(level=logging.DEBUG)
- self.logger.info("Starting the scheduler")
- dagbag = models.DagBag(self.subdir, sync_to_db=True)
- executor = self.executor = dagbag.executor
- executor.start()
+ # DAGs can be pickled for easier remote execution by some executors
+ pickle_dags = False
+ if self.do_pickle and self.executor.__class__ not in \
+ (executors.LocalExecutor, executors.SequentialExecutor):
+ pickle_dags = True
+
+ # Use multiple processes to parse and generate tasks for the
+ # DAGs in parallel. By processing them in separate processes,
+ # we can get parallelism and isolation from potentially harmful
+ # user code.
+ self.logger.info("Processing files using up to {} processes at a time "
+ .format(self.max_threads))
+ self.logger.info("Running execute loop for {} seconds"
+ .format(self.run_duration))
+ self.logger.info("Processing each file at most {} times"
+ .format(self.num_runs))
+ self.logger.info("Process each file at most once every {} seconds"
+ .format(self.file_process_interval))
+ self.logger.info("Checking for new files in {} every {} seconds"
+ .format(self.subdir, self.dag_dir_list_interval))
+
+ # Build up a list of Python files that could contain DAGs
+ self.logger.info("Searching for files in {}".format(self.subdir))
+ known_file_paths = list_py_file_paths(self.subdir)
+ self.logger.info("There are {} files in {}"
+ .format(len(known_file_paths), self.subdir))
+
+ def processor_factory(file_path, log_file_path):
+ return DagFileProcessor(file_path,
+ pickle_dags,
+ self.dag_ids,
+ log_file_path)
+
+ processor_manager = DagFileProcessorManager(self.subdir,
+ known_file_paths,
+ self.max_threads,
+ self.file_process_interval,
+ self.child_process_log_directory,
+ self.num_runs,
+ processor_factory)
+ try:
+ self._execute_helper(processor_manager)
+ finally:
+ self.logger.info("Exited execute loop")
+
+ # Kill all child processes on exit since we don't want to leave
+ # them as orphaned.
+ pids_to_kill = processor_manager.get_all_pids()
+ if len(pids_to_kill) > 0:
+ # First try SIGTERM
+ this_process = psutil.Process(os.getpid())
+ # Only check child processes to ensure that we don't have a case
+ # where we kill the wrong process because a child process died
+ # but the PID got reused.
+ child_processes = [x for x in this_process.children(recursive=True)
+ if x.is_running() and x.pid in pids_to_kill]
+ for child in child_processes:
+ self.logger.info("Terminating child PID: {}".format(child.pid))
+ child.terminate()
+ timeout = 5
+ self.logger.info("Waiting up to {}s for processes to exit..."
+ .format(timeout))
+ try:
+ psutil.wait_procs(child_processes, timeout)
+ except psutil.TimeoutExpired:
+ self.logger.debug("Ran out of time while waiting for "
+ "processes to exit")
+
+ # Then SIGKILL
+ child_processes = [x for x in this_process.children(recursive=True)
+ if x.is_running() and x.pid in pids_to_kill]
+ if len(child_processes) > 0:
+ for child in child_processes:
+ self.logger.info("Killing child PID: {}".format(child.pid))
+ child.kill()
+ child.wait()
+
+ def _execute_helper(self, processor_manager):
+ """
+ :param processor_manager: manager to use
+ :type processor_manager: DagFileProcessorManager
+ :return: None
+ """
+ self.executor.start()
+
+ session = settings.Session()
+ self.logger.info("Resetting state for orphaned tasks")
# grab orphaned tasks and make sure to reset their state
active_runs = DagRun.find(
state=State.RUNNING,
@@ -708,103 +1279,250 @@ class SchedulerJob(BaseJob):
session=session
)
for dr in active_runs:
+ self.logger.info("Resetting {} {}".format(dr.dag_id,
+ dr.execution_date))
self._reset_state_for_orphaned_tasks(dr, session=session)
- self.runs = 0
- while not self.num_runs or self.num_runs > self.runs:
- try:
- loop_start_dttm = datetime.now()
- try:
- self.prioritize_queued(executor=executor, dagbag=dagbag)
- except Exception as e:
- self.logger.exception(e)
+ self.logger.info("Removing old import errors")
+ self.clear_import_errors(session)
+ session.close()
- self.runs += 1
- try:
- if self.runs % self.refresh_dags_every == 0:
- dagbag = models.DagBag(self.subdir, sync_to_db=True)
- else:
- dagbag.collect_dags(only_if_updated=True)
- except Exception as e:
- self.logger.error("Failed at reloading the dagbag. {}".format(e))
- Stats.incr('dag_refresh_error', 1, 1)
- sleep(5)
-
- if len(self.dag_ids) > 0:
- dags = [dag for dag in dagbag.dags.values() if dag.dag_id in self.dag_ids]
- else:
- dags = [
- dag for dag in dagbag.dags.values()
- if not dag.parent_dag]
-
- paused_dag_ids = dagbag.paused_dags()
- dags = [x for x in dags if x.dag_id not in paused_dag_ids]
- # dags = filter(lambda x: x.dag_id not in paused_dag_ids, dags)
-
- self.logger.debug("Total Cores: {} Max Threads: {} DAGs:{}".
- format(multiprocessing.cpu_count(),
- self.max_threads,
- len(dags)))
- dags = self._split(dags, math.ceil(len(dags) / self.max_threads))
- tis_q = multiprocessing.Queue()
- jobs = [multiprocessing.Process(target=self._do_dags,
- args=(dagbag, dags[i], tis_q))
- for i in range(len(dags))]
-
- self.logger.info("Starting {} scheduler jobs".format(len(jobs)))
- for j in jobs:
- j.start()
-
- while any(j.is_alive() for j in jobs):
- while not tis_q.empty():
- ti_key, pickle_id = tis_q.get()
- dag = dagbag.dags[ti_key[0]]
- task = dag.get_task(ti_key[1])
- ti = TI(task, ti_key[2])
- ti.refresh_from_db(session=session, lock_for_update=True)
- if ti.state == State.SCHEDULED:
- session.commit()
- self.logger.debug("Task {} was picked up by another scheduler"
- .format(ti))
- continue
- elif ti.state is State.NONE:
- ti.state = State.SCHEDULED
-
- self.executor.queue_task_instance(ti, pickle_id=pickle_id)
+ execute_start_time = datetime.now()
+
+ # Last time stats were printed
+ last_stat_print_time = datetime(2000, 1, 1)
+ # Last time that self.heartbeat() was called.
+ last_self_heartbeat_time = datetime.now()
+ # Last time that the DAG dir was traversed to look for files
+ last_dag_dir_refresh_time = datetime.now()
+
+ # Use this value initially
+ known_file_paths = processor_manager.file_paths
+
+ # For the execute duration, parse and schedule DAGs
+ while (datetime.now() - execute_start_time).total_seconds() < \
+ self.run_duration:
+ self.logger.debug("Starting Loop...")
+ loop_start_time = time.time()
+
+ # Traverse the DAG directory for Python files containing DAGs
+ # periodically
+ elapsed_time_since_refresh = (datetime.now() -
+ last_dag_dir_refresh_time).total_seconds()
+
+ if elapsed_time_since_refresh > self.dag_dir_list_interval:
+ # Build up a list of Python files that could contain DAGs
+ self.logger.info("Searching for files in {}".format(self.subdir))
+ known_file_paths = list_py_file_paths(self.subdir)
+ last_dag_dir_refresh_time = datetime.now()
+ self.logger.info("There are {} files in {}"
+ .format(len(known_file_paths), self.subdir))
+ processor_manager.set_file_paths(known_file_paths)
+
+ # Kick of new processes and collect results from finished ones
+ self.logger.info("Heartbeating the process manager")
+ simple_dags = processor_manager.heartbeat()
+
+ if self.using_sqlite:
+ # For the sqlite case w/ 1 thread, wait until the processor
+ # is finished to avoid concurrent access to the DB.
+ self.logger.debug("Waiting for processors to finish since we're "
+ "using sqlite")
+ processor_manager.wait_until_finished()
+
+ # Send tasks for execution if available
+ if len(simple_dags) > 0:
+ simple_dag_bag = SimpleDagBag(simple_dags)
+
+ # Handle cases where a DAG run state is set (perhaps manually) to
+ # a non-running state. Handle task instances that belong to
+ # DAG runs in those states
+
+ # If a task instance is up for retry but the corresponding DAG run
+ # isn't running, mark the task instance as FAILED so we don't try
+ # to re-run it.
+ self._change_state_for_tis_without_dagrun(simple_dag_bag,
+ [State.UP_FOR_RETRY],
+ State.FAILED)
+ # If a task instance is scheduled or queued, but the corresponding
+ # DAG run isn't running, set the state to NONE so we don't try to
+ # re-run it.
+ self._change_state_for_tis_without_dagrun(simple_dag_bag,
+ [State.QUEUED,
+ State.SCHEDULED],
+ State.NONE)
+
+ self._execute_task_instances(simple_dag_bag,
+ (State.SCHEDULED,
+ State.UP_FOR_RETRY))
+
+ # Call hearbeats
+ self.logger.info("Heartbeating the executor")
+ self.executor.heartbeat()
+
+ # Process events from the executor
+ self._process_executor_events()
+
+ # Heartbeat the scheduler periodically
+ time_since_last_heartbeat = (datetime.now() -
+ last_self_heartbeat_time).total_seconds()
+ if time_since_last_heartbeat > self.heartrate:
+ self.logger.info("Heartbeating the scheduler")
+ self.heartbeat()
+ last_self_heartbeat_time = datetime.now()
+
+ # Occasionally print out stats about how fast the files are getting processed
+ if ((datetime.now() - last_stat_print_time).total_seconds() >
+ self.print_stats_interval):
+ if len(known_file_paths) > 0:
+ self._log_file_processing_stats(known_file_paths,
+ processor_manager)
+ last_stat_print_time = datetime.now()
+
+ loop_end_time = time.time()
+ self.logger.debug("Ran scheduling loop in {:.2f}s"
+ .format(loop_end_time - loop_start_time))
+ self.logger.debug("Sleeping for {:.2f}s"
+ .format(self._processor_poll_interval))
+ time.sleep(self._processor_poll_interval)
+
+ # Exit early for a test mode
+ if processor_manager.max_runs_reached():
+ self.logger.info("Exiting loop as all files have been processed "
+ "{} times".format(self.num_runs))
+ break
+
+ # Stop any processors
+ processor_manager.terminate()
+
+ # Verify that all files were processed, and if so, deactivate DAGs that
+ # haven't been touched by the scheduler as they likely have been
+ # deleted.
+ all_files_processed = True
+ for file_path in known_file_paths:
+ if processor_manager.get_last_finish_time(file_path) is None:
+ all_files_processed = False
+ break
+ if all_files_processed:
+ self.logger.info("Deactivating DAGs that haven't been touched since {}"
+ .format(execute_start_time.isoformat()))
+ models.DAG.deactivate_stale_dags(execute_start_time)
+
+ self.executor.end()
+
+ settings.Session.remove()
- session.merge(ti)
- session.commit()
+ @provide_session
+ def process_file(self, file_path, pickle_dags=False, session=None):
+ """
+ Process a Python file containing Airflow DAGs.
+
+ This includes:
+
+ 1. Execute the file and look for DAG objects in the namespace.
+ 2. Pickle the DAG and save it to the DB (if necessary).
+ 3. For each DAG, see what tasks should run and create appropriate task
+ instances in the DB.
+ 4. Record any errors importing the file into ORM
+ 5. Kill (in ORM) any task instances belonging to the DAGs that haven't
+ issued a heartbeat in a while.
+
+ Returns a list of SimpleDag objects that represent the DAGs found in
+ the file
+
+ :param file_path: the path to the Python file that should be executed
+ :type file_path: unicode
+ :param pickle_dags: whether serialize the DAGs found in the file and
+ save them to the db
+ :type pickle_dags: bool
+ :return: a list of SimpleDags made from the Dags found in the file
+ :rtype: list[SimpleDag]
+ """
+ self.logger.info("Processing file {} for tasks to queue".format(file_path))
+ # As DAGs are parsed from this file, they will be converted into SimpleDags
+ simple_dags = []
+
+ try:
+ dagbag = models.DagBag(file_path)
+ except Exception:
+ self.logger.exception("Failed at reloading the DAG file {}".format(file_path))
+ Stats.incr('dag_file_refresh_error', 1, 1)
+ return []
+
+ if len(dagbag.dags) > 0:
+ self.logger.info("DAG(s) {} retrieved from {}"
+ .format(dagbag.dags.keys(),
+ file_path))
+ else:
+ self.logger.warn("No viable dags retrieved from {}".format(file_path))
+ return []
+
+ # Save individual DAGs in the ORM and update DagModel.last_scheduled_time
+ sync_time = datetime.now()
+ for dag in dagbag.dags.values():
+ models.DAG.sync_to_db(dag, dag.owner, sync_time)
+
+ paused_dag_ids = [dag.dag_id for dag in dagbag.dags.values()
+ if dag.is_paused]
+
+ # Pickle the DAGs (if necessary) and put them into a SimpleDag
+ for dag_id in dagbag.dags:
+ dag = dagbag.get_dag(dag_id)
+ pickle_id = None
+ if pickle_dags:
+ pickle_id = dag.pickle(session).id
+
+ task_ids = [task.task_id for task in dag.tasks]
+
+ # Only return DAGs that are not paused
+ if dag_id not in paused_dag_ids:
+ simple_dags.append(SimpleDag(dag.dag_id,
+ task_ids,
+ dag.full_filepath,
+ dag.concurrency,
+ dag.is_paused,
+ pickle_id))
+
+ if len(self.dag_ids) > 0:
+ dags = [dag for dag in dagbag.dags.values()
+ if dag.dag_id in self.dag_ids and
+ dag.dag_id not in paused_dag_ids]
+ else:
+ dags = [dag for dag in dagbag.dags.values()
+ if not dag.parent_dag and
+ dag.dag_id not in paused_dag_ids]
+
+ # Not using multiprocessing.Queue() since it's no longer a separate
+ # process and due to some unusual behavior. (empty() incorrectly
+ # returns true?)
+ ti_keys_to_schedule = []
+
+ self._process_dags(dagbag, dags, ti_keys_to_schedule)
+
+ for ti_key in ti_keys_to_schedule:
+ dag = dagbag.dags[ti_key[0]]
+ task = dag.get_task(ti_key[1])
+ ti = models.TaskInstance(task, ti_key[2])
+ # Task starts out in the scheduled state. All tasks in the
+ # scheduled state will be sent to the executor
+ ti.state = State.SCHEDULED
+
+ # Also save this task instance to the DB.
+ self.logger.info("Creating / updating {} in ORM".format(ti))
+ session.merge(ti)
+ session.commit()
- for j in jobs:
- j.join()
+ # Record import errors into the ORM
+ try:
+ self.record_import_errors(session, dagbag)
+ except Exception:
+ self.logger.exception("Error logging import errors!")
+ try:
+ dagbag.kill_zombies()
+ except Exception:
+ self.logger.exception("Error killing zombies!")
- self.logger.info("Done queuing tasks, calling the executor's "
- "heartbeat")
- duration_sec = (datetime.now() - loop_start_dttm).total_seconds()
- self.logger.info("Loop took: {} seconds".format(duration_sec))
- Stats.timing("scheduler_loop", duration_sec * 1000)
- try:
- self.import_errors(dagbag)
- except Exception as e:
- self.logger.exception(e)
- try:
- dagbag.kill_zombies()
- except Exception as e:
- self.logger.exception(e)
- try:
- # We really just want the scheduler to never ever stop.
- executor.heartbeat()
- self.heartbeat()
- except Exception as e:
- self.logger.exception(e)
- self.logger.error("Tachycardia!")
- except Exception as deep_e:
- self.logger.exception(deep_e)
- raise
- finally:
- settings.Session.remove()
- executor.end()
- session.close()
+ return simple_dags
@provide_session
def heartbeat_callback(self, session=None):
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/models.py
----------------------------------------------------------------------
diff --git a/airflow/models.py b/airflow/models.py
index f589b3e..8326a92 100644
--- a/airflow/models.py
+++ b/airflow/models.py
@@ -59,6 +59,7 @@ from airflow import settings, utils
from airflow.executors import DEFAULT_EXECUTOR, LocalExecutor
from airflow import configuration
from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.dag.base_dag import BaseDag, BaseDagBag
from airflow.utils.dates import cron_presets, date_range as utils_date_range
from airflow.utils.db import provide_session
from airflow.utils.decorators import apply_defaults
@@ -129,7 +130,7 @@ def clear_task_instances(tis, session, activate_dag_runs=True):
dr.start_date = datetime.now()
-class DagBag(LoggingMixin):
+class DagBag(BaseDagBag, LoggingMixin):
"""
A dagbag is a collection of dags, parsed out of a folder tree and has high
level configuration settings, like what database to use as a backend and
@@ -140,7 +141,7 @@ class DagBag(LoggingMixin):
independent settings sets.
:param dag_folder: the folder to scan to find DAGs
- :type dag_folder: str
+ :type dag_folder: unicode
:param executor: the executor to use when executing task instances
in this DagBag
:param include_examples: whether to include the examples that ship
@@ -155,26 +156,23 @@ class DagBag(LoggingMixin):
self,
dag_folder=None,
executor=DEFAULT_EXECUTOR,
- include_examples=configuration.getboolean('core', 'LOAD_EXAMPLES'),
- sync_to_db=False):
+ include_examples=configuration.getboolean('core', 'LOAD_EXAMPLES')):
dag_folder = dag_folder or DAGS_FOLDER
self.logger.info("Filling up the DagBag from {}".format(dag_folder))
self.dag_folder = dag_folder
self.dags = {}
- self.sync_to_db = sync_to_db
# the file's last modified timestamp when we last read it
self.file_last_changed = {}
self.executor = executor
self.import_errors = {}
+
if include_examples:
example_dag_folder = os.path.join(
os.path.dirname(__file__),
'example_dags')
self.collect_dags(example_dag_folder)
self.collect_dags(dag_folder)
- if sync_to_db:
- self.deactivate_inactive_dags()
def size(self):
"""
@@ -307,7 +305,7 @@ class DagBag(LoggingMixin):
return found_dags
@provide_session
- def kill_zombies(self, session):
+ def kill_zombies(self, session=None):
"""
Fails tasks that haven't had a heartbeat in too long
"""
@@ -355,20 +353,6 @@ class DagBag(LoggingMixin):
for task in dag.tasks:
settings.policy(task)
- if self.sync_to_db:
- session = settings.Session()
- orm_dag = session.query(
- DagModel).filter(DagModel.dag_id == dag.dag_id).first()
- if not orm_dag:
- orm_dag = DagModel(dag_id=dag.dag_id)
- orm_dag.fileloc = root_dag.full_filepath
- orm_dag.is_subdag = dag.is_subdag
- orm_dag.owners = root_dag.owner
- orm_dag.is_active = True
- session.merge(orm_dag)
- session.commit()
- session.close()
-
for subdag in dag.subdags:
subdag.full_filepath = dag.full_filepath
subdag.parent_dag = dag
@@ -756,8 +740,77 @@ class TaskInstance(Base):
the orchestrator.
"""
dag = self.task.dag
- iso = self.execution_date.isoformat()
- cmd = "airflow run {self.dag_id} {self.task_id} {iso} "
+
+ # Keeping existing logic, but not entirely sure why this is here.
+ if not pickle_id and dag:
+ if dag.full_filepath != dag.filepath:
+ path = "DAGS_FOLDER/{}".format(dag.filepath)
+ elif dag.full_filepath:
+ path = dag.full_filepath
+
+ return TaskInstance.generate_command(
+ self.dag_id,
+ self.task_id,
+ self.execution_date,
+ mark_success=mark_success,
+ ignore_dependencies=ignore_dependencies,
+ ignore_depends_on_past=ignore_depends_on_past,
+ force=force,
+ local=local,
+ pickle_id=pickle_id,
+ file_path=path,
+ raw=raw,
+ job_id=job_id,
+ pool=pool)
+
+ @staticmethod
+ def generate_command(dag_id,
+ task_id,
+ execution_date,
+ mark_success=False,
+ ignore_dependencies=False,
+ ignore_depends_on_past=False,
+ force=False,
+ local=False,
+ pickle_id=None,
+ file_path=None,
+ raw=False,
+ job_id=None,
+ pool=None
+ ):
+ """
+ Generates the shell command required to execute this task instance.
+
+ :param dag_id: DAG ID
+ :type dag_id: unicode
+ :param task_id: Task ID
+ :type task_id: unicode
+ :param execution_date: Execution date for the task
+ :type execution_date: datetime
+ :param mark_success: Whether to mark the task as successful
+ :type mark_success: bool
+ :param ignore_dependencies: Whether to ignore the dependencies and run
+ anyway
+ :type ignore_dependencies: bool
+ :param ignore_depends_on_past: Whether to ignore the depends on past
+ setting and run anyway
+ :type ignore_depends_on_past: bool
+ :param force: Whether to force running - see TaskInstance.run()
+ :type force: bool
+ :param local: Whether to run the task locally
+ :type local: bool
+ :param pickle_id: If the DAG was serialized to the DB, the ID
+ associated with the pickled DAG
+ :type pickle_id: unicode
+ :param file_path: path to the file containing the DAG definition
+ :param raw: raw mode (needs more details)
+ :param job_id: job ID (needs more details)
+ :param pool: the Airflow pool that the task should run in
+ :type pool: unicode
+ :return: shell command that can be used to run the task instance
+ """
+ iso = execution_date.isoformat()
+ cmd = "airflow run {dag_id} {task_id} {iso} "
cmd += "--mark_success " if mark_success else ""
cmd += "--pickle {pickle_id} " if pickle_id else ""
cmd += "--job_id {job_id} " if job_id else ""
@@ -767,11 +820,7 @@ class TaskInstance(Base):
cmd += "--local " if local else ""
cmd += "--pool {pool} " if pool else ""
cmd += "--raw " if raw else ""
- if not pickle_id and dag:
- if dag.full_filepath != dag.filepath:
- cmd += "-sd DAGS_FOLDER/{dag.filepath} "
- elif dag.full_filepath:
- cmd += "-sd {dag.full_filepath}"
+ cmd += "-sd {file_path}"
return cmd.format(**locals())
@property
@@ -1185,9 +1234,7 @@ class TaskInstance(Base):
.first()
)
if not pool:
- raise ValueError(
- "Task specified a pool ({}) but the pool "
- "doesn't exist!".format(self.task.pool))
+ return False
open_slots = pool.open_slots(session=session)
return open_slots <= 0
@@ -2469,7 +2516,7 @@ class DagModel(Base):
@functools.total_ordering
-class DAG(LoggingMixin):
+class DAG(BaseDag, LoggingMixin):
"""
A dag (directed acyclic graph) is a collection of tasks with directional
dependencies. A dag also has a schedule, a start end an end date
@@ -2557,8 +2604,14 @@ class DAG(LoggingMixin):
del self.default_args['params']
validate_key(dag_id)
+
+ # Properties from BaseDag
+ self._dag_id = dag_id
+ self._full_filepath = full_filepath if full_filepath else ''
+ self._concurrency = concurrency
+ self._pickle_id = None
+
self.task_dict = dict()
- self.dag_id = dag_id
self.start_date = start_date
self.end_date = end_date
self.schedule_interval = schedule_interval
@@ -2568,14 +2621,12 @@ class DAG(LoggingMixin):
self._schedule_interval = None
else:
self._schedule_interval = schedule_interval
- self.full_filepath = full_filepath if full_filepath else ''
if isinstance(template_searchpath, six.string_types):
template_searchpath = [template_searchpath]
self.template_searchpath = template_searchpath
self.parent_dag = None # Gets set when DAGs are loaded
self.last_loaded = datetime.now()
self.safe_dag_id = dag_id.replace('.', '__dot__')
- self.concurrency = concurrency
self.max_active_runs = max_active_runs
self.dagrun_timeout = dagrun_timeout
self.sla_miss_callback = sla_miss_callback
@@ -2597,7 +2648,9 @@ class DAG(LoggingMixin):
def __eq__(self, other):
return (
type(self) == type(other) and
- all(self.__dict__.get(c, None) == other.__dict__.get(c, None)
+ # Use getattr() instead of __dict__ as __dict__ doesn't return
+ # correct values for properties.
+ all(getattr(self, c, None) == getattr(other, c, None)
for c in self._comps))
def __ne__(self, other):
@@ -2672,6 +2725,38 @@ class DAG(LoggingMixin):
return dttm
@property
+ def dag_id(self):
+ return self._dag_id
+
+ @dag_id.setter
+ def dag_id(self, value):
+ self._dag_id = value
+
+ @property
+ def full_filepath(self):
+ return self._full_filepath
+
+ @full_filepath.setter
+ def full_filepath(self, value):
+ self._full_filepath = value
+
+ @property
+ def concurrency(self):
+ return self._concurrency
+
+ @concurrency.setter
+ def concurrency(self, value):
+ self._concurrency = value
+
+ @property
+ def pickle_id(self):
+ return self._pickle_id
+
+ @pickle_id.setter
+ def pickle_id(self, value):
+ self._pickle_id = value
+
+ @property
def tasks(self):
return list(self.task_dict.values())
@@ -3134,6 +3219,80 @@ class DAG(LoggingMixin):
run.refresh_from_db()
return run
+ @staticmethod
+ @provide_session
+ def sync_to_db(dag, owner, sync_time, session=None):
+ """
+ Save attributes about this DAG to the DB. Note that this method
+ can be called for both DAGs and SubDAGs. A SubDag is actually a
+ SubDagOperator.
+
+ :param dag: the DAG object to save to the DB
+ :type dag: DAG
+ :own
+ :param sync_time: The time that the DAG should be marked as sync'ed
+ :type sync_time: datetime
+ :return: None
+ """
+ orm_dag = session.query(
+ DagModel).filter(DagModel.dag_id == dag.dag_id).first()
+ if not orm_dag:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ logging.info("Creating ORM DAG for %s",
+ dag.dag_id)
+ orm_dag.fileloc = dag.full_filepath
+ orm_dag.is_subdag = dag.is_subdag
+ orm_dag.owners = owner
+ orm_dag.is_active = True
+ orm_dag.last_scheduler_run = sync_time
+ session.merge(orm_dag)
+ session.commit()
+
+ for subdag in dag.subdags:
+ DAG.sync_to_db(subdag, owner, sync_time, session=session)
+
+ @staticmethod
+ @provide_session
+ def deactivate_unknown_dags(active_dag_ids, session=None):
+ """
+ Given a list of known DAGs, deactivate any other DAGs that are
+ marked as active in the ORM
+
+ :param active_dag_ids: list of DAG IDs that are active
+ :type active_dag_ids: list[unicode]
+ :return: None
+ """
+
+ if len(active_dag_ids) == 0:
+ return
+ for dag in session.query(
+ DagModel).filter(~DagModel.dag_id.in_(active_dag_ids)).all():
+ dag.is_active = False
+ session.merge(dag)
+
+ @staticmethod
+ @provide_session
+ def deactivate_stale_dags(expiration_date, session=None):
+ """
+ Deactivate any DAGs that were last touched by the scheduler before
+ the expiration date. These DAGs were likely deleted.
+
+ :param expiration_date: set inactive DAGs that were touched before this
+ time
+ :type expiration_date: datetime
+ :return: None
+ """
+ for dag in session.query(
+ DagModel).filter(DagModel.last_scheduler_run < expiration_date,
+ DagModel.is_active).all():
+ logging.info("Deactivating DAG ID %s since it was last touched "
+ "by the scheduler at %s",
+ dag.dag_id,
+ dag.last_scheduler_run.isoformat())
+ dag.is_active = False
+ session.merge(dag)
+ session.commit()
+
class Chart(Base):
__tablename__ = "chart"
@@ -3567,7 +3726,7 @@ class DagRun(Base):
state=State.unfinished(),
session=session
)
- none_depends_on_past = all(t.task.depends_on_past for t in unfinished_tasks)
+ none_depends_on_past = all(not t.task.depends_on_past for t in unfinished_tasks)
# small speed up
if unfinished_tasks and none_depends_on_past:
@@ -3642,6 +3801,44 @@ class DagRun(Base):
session.commit()
+ @staticmethod
+ def get_running_tasks(session, dag_id, task_ids):
+ """
+ Returns the number of tasks running in the given DAG.
+
+ :param session: ORM session
+ :param dag_id: ID of the DAG to get the task concurrency of
+ :type dag_id: unicode
+ :param task_ids: A list of valid task IDs for the given DAG
+ :type task_ids: list[unicode]
+ :return: The number of running tasks
+ :rtype: int
+ """
+ qry = session.query(func.count(TaskInstance.task_id)).filter(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.task_id.in_(task_ids),
+ TaskInstance.state == State.RUNNING,
+ )
+ return qry.scalar()
+
+ @staticmethod
+ def get_run(session, dag_id, execution_date):
+ """
+ :param dag_id: DAG ID
+ :type dag_id: unicode
+ :param execution_date: execution date
+ :type execution_date: datetime
+ :return: DagRun corresponding to the given dag_id and execution date
+ if one exists. None otherwise.
+ :rtype: DagRun
+ """
+ qry = session.query(DagRun).filter(
+ DagRun.dag_id == dag_id,
+ DagRun.external_trigger == False,
+ DagRun.execution_date == execution_date,
+ )
+ return qry.first()
+
class Pool(Base):
__tablename__ = "slot_pool"
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fdb7e949/airflow/settings.py
----------------------------------------------------------------------
diff --git a/airflow/settings.py b/airflow/settings.py
index 9f9bb14..ccd77ee 100644
--- a/airflow/settings.py
+++ b/airflow/settings.py
@@ -72,20 +72,11 @@ SQL_ALCHEMY_CONN = conf.get('core', 'SQL_ALCHEMY_CONN')
LOGGING_LEVEL = logging.INFO
DAGS_FOLDER = os.path.expanduser(conf.get('core', 'DAGS_FOLDER'))
-engine_args = {}
-if 'sqlite' not in SQL_ALCHEMY_CONN:
- # Engine args not supported by sqlite
- engine_args['pool_size'] = conf.getint('core', 'SQL_ALCHEMY_POOL_SIZE')
- engine_args['pool_recycle'] = conf.getint('core',
- 'SQL_ALCHEMY_POOL_RECYCLE')
-
-engine = create_engine(SQL_ALCHEMY_CONN, **engine_args)
-Session = scoped_session(
- sessionmaker(autocommit=False, autoflush=False, bind=engine))
-
# can't move this to conf due to ConfigParser interpolation
LOG_FORMAT = (
'[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s')
+LOG_FORMAT_WITH_THREAD_NAME = (
+ '[%(asctime)s] {%(filename)s:%(lineno)d} %(threadName)s %(levelname)s - %(message)s')
SIMPLE_LOG_FORMAT = '%(asctime)s %(levelname)s - %(message)s'
@@ -116,10 +107,28 @@ def policy(task_instance):
pass
-def configure_logging():
+def configure_logging(log_format=LOG_FORMAT):
logging.root.handlers = []
logging.basicConfig(
- format=LOG_FORMAT, stream=sys.stdout, level=LOGGING_LEVEL)
+ format=log_format, stream=sys.stdout, level=LOGGING_LEVEL)
+
+engine = None
+Session = None
+
+
+def configure_orm():
+ global engine
+ global Session
+ engine_args = {}
+ if 'sqlite' not in SQL_ALCHEMY_CONN:
+ # Engine args not supported by sqlite
+ engine_args['pool_size'] = conf.getint('core', 'SQL_ALCHEMY_POOL_SIZE')
+ engine_args['pool_recycle'] = conf.getint('core',
+ 'SQL_ALCHEMY_POOL_RECYCLE')
+
+ engine = create_engine(SQL_ALCHEMY_CONN, **engine_args)
+ Session = scoped_session(
+ sessionmaker(autocommit=False, autoflush=False, bind=engine))
try:
from airflow_local_settings import *
@@ -128,3 +137,4 @@ except:
pass
configure_logging()
+configure_orm()