airflow-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From davy...@apache.org
Subject incubator-airflow git commit: [AIRFLOW-680] Disable connection pool for commands
Date Mon, 12 Dec 2016 23:32:43 GMT
Repository: incubator-airflow
Updated Branches:
  refs/heads/master 15ff540ec -> 67ab416db


[AIRFLOW-680] Disable connection pool for commands

This is a continuation of Max's PR here:
https://github.com/apache/incubator-
airflow/pull/1021
We have seen a very substantial DB cpu usage
decrease from this PR (~10x).

This PR was originally created by plypaul I am
just cherrypicking onto master for him.

Testing Done:
- Has been running in Airbnb production for quite
some time (off of a different merge base though)

Closes #1925 from aoen/ddavydov/reduce_db_sessions


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/67ab416d
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/67ab416d
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/67ab416d

Branch: refs/heads/master
Commit: 67ab416dbc6f61095e6e102f44b3825fc482b6c5
Parents: 15ff540
Author: Paul Yang <paul.yang@airbnb.com>
Authored: Mon Dec 12 15:32:25 2016 -0800
Committer: Dan Davydov <dan.davydov@airbnb.com>
Committed: Mon Dec 12 15:32:29 2016 -0800

----------------------------------------------------------------------
 airflow/bin/cli.py  |  3 +++
 airflow/jobs.py     | 22 ++++++++++++++--------
 airflow/settings.py |  7 +++++--
 3 files changed, 22 insertions(+), 10 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/67ab416d/airflow/bin/cli.py
----------------------------------------------------------------------
diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py
index 21e1d23..eab4a30 100755
--- a/airflow/bin/cli.py
+++ b/airflow/bin/cli.py
@@ -321,6 +321,9 @@ def set_is_paused(is_paused, args, dag=None):
 
 
 def run(args, dag=None):
+    # Disable connection pooling to reduce the # of connections on the DB
+    # while it's waiting for the task to finish.
+    settings.configure_orm(disable_connection_pool=True)
     db_utils.pessimistic_connection_handling()
     if dag:
         args.dag_id = dag.dag_id

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/67ab416d/airflow/jobs.py
----------------------------------------------------------------------
diff --git a/airflow/jobs.py b/airflow/jobs.py
index 229424e..a2d94e3 100644
--- a/airflow/jobs.py
+++ b/airflow/jobs.py
@@ -165,15 +165,21 @@ class BaseJob(Base, LoggingMixin):
         if job.state == State.SHUTDOWN:
             self.kill()
 
+        # Figure out how long to sleep for
+        sleep_for = 0
         if job.latest_heartbeat:
-            sleep_for = self.heartrate - (
-                datetime.now() - job.latest_heartbeat).total_seconds()
-            if sleep_for > 0:
-                sleep(sleep_for)
+            sleep_for = max(
+                0,
+                self.heartrate - (datetime.now() - job.latest_heartbeat).total_seconds())
 
-        job.latest_heartbeat = datetime.now()
+        # Don't keep session open while sleeping as it leaves a connection open
+        session.close()
+        sleep(sleep_for)
 
+        # Update last heartbeat time
         session = settings.Session()
+        job = session.query(BaseJob).filter(BaseJob.id == self.id).first()
+        job.latest_heartbeat = datetime.now()
         session.merge(job)
         session.commit()
 
@@ -182,7 +188,7 @@ class BaseJob(Base, LoggingMixin):
         self.logger.debug('[heart] Boom.')
 
     def run(self):
-        Stats.incr(self.__class__.__name__.lower()+'_start', 1, 1)
+        Stats.incr(self.__class__.__name__.lower() + '_start', 1, 1)
         # Adding an entry in the DB
         session = settings.Session()
         self.state = State.RUNNING
@@ -202,7 +208,7 @@ class BaseJob(Base, LoggingMixin):
         session.commit()
         session.close()
 
-        Stats.incr(self.__class__.__name__.lower()+'_end', 1, 1)
+        Stats.incr(self.__class__.__name__.lower() + '_end', 1, 1)
 
     def _execute(self):
         raise NotImplementedError("This method needs to be overridden")
@@ -714,7 +720,7 @@ class SchedulerJob(BaseJob):
                 .filter(or_(
                     DagRun.external_trigger == False,
                     # add % as a wildcard for the like query
-                    DagRun.run_id.like(DagRun.ID_PREFIX+'%')
+                    DagRun.run_id.like(DagRun.ID_PREFIX + '%')
                 ))
             )
             last_scheduled_run = qry.scalar()

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/67ab416d/airflow/settings.py
----------------------------------------------------------------------
diff --git a/airflow/settings.py b/airflow/settings.py
index 60efb4a..e8da674 100644
--- a/airflow/settings.py
+++ b/airflow/settings.py
@@ -23,6 +23,7 @@ import sys
 
 from sqlalchemy import create_engine
 from sqlalchemy.orm import scoped_session, sessionmaker
+from sqlalchemy.pool import NullPool
 
 from airflow import configuration as conf
 
@@ -121,11 +122,13 @@ engine = None
 Session = None
 
 
-def configure_orm():
+def configure_orm(disable_connection_pool=False):
     global engine
     global Session
     engine_args = {}
-    if 'sqlite' not in SQL_ALCHEMY_CONN:
+    if disable_connection_pool:
+        engine_args['poolclass'] = NullPool
+    elif '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',


Mime
View raw message