[AIRFLOW-4956] Fix LocalTaskJob heartbeat log spamming

airflow/jobs/base_job.py

@@ -35,6 +35,7 @@
 from airflow.stats import Stats
 from airflow.utils import helpers, timezone
 from airflow.utils.db import create_session, provide_session
+from airflow.utils.helpers import convert_camel_to_snake
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.net import get_hostname
 from airflow.utils.sqlalchemy import UtcDateTime
@@ -191,8 +192,11 @@ def heartbeat(self):
 
                 self.heartbeat_callback(session=session)
                 self.log.debug('[heartbeat]')
-        except OperationalError as e:
-            self.log.error("Scheduler heartbeat got an exception: %s", str(e))
+        except OperationalError:
+            Stats.incr(
+                convert_camel_to_snake(self.__class__.__name__) + '_heartbeat_failure', 1,
+                1)
+            self.log.exception("%s heartbeat got an exception", self.__class__.__name__)
 
     def run(self):
         Stats.incr(self.__class__.__name__.lower() + '_start', 1, 1)

airflow/jobs/local_task_job.py

@@ -22,8 +22,6 @@
 import signal
 import time
 
-from sqlalchemy.exc import OperationalError
-
 from airflow import configuration as conf
 from airflow.exceptions import AirflowException
 from airflow.stats import Stats
@@ -91,7 +89,6 @@ def signal_handler(signum, frame):
         try:
             self.task_runner.start()
 
-            last_heartbeat_time = time.time()
             heartbeat_time_limit = conf.getint('scheduler',
                                                'scheduler_zombie_task_threshold')
             while True:
@@ -103,16 +100,8 @@ def signal_handler(signum, frame):
 
                 # Periodically heartbeat so that the scheduler doesn't think this
                 # is a zombie
-                try:
-                    self.heartbeat()
-                    last_heartbeat_time = time.time()
-                except OperationalError:
-                    Stats.incr('local_task_job_heartbeat_failure', 1, 1)
-                    self.log.exception(
-                        "Exception while trying to heartbeat! Sleeping for %s seconds",
-                        self.heartrate
-                    )
-                    time.sleep(self.heartrate)
+                last_heartbeat_time = time.time()
+                self.heartbeat()
 
                 # If it's been too long since we've heartbeat, then it's possible that
                 # the scheduler rescheduled this task, so kill launched processes.
@@ -124,6 +113,13 @@ def signal_handler(signum, frame):
                                            "exceeded limit ({}s)."
                                            .format(time_since_last_heartbeat,
                                                    heartbeat_time_limit))
+
+                if time_since_last_heartbeat < self.heartrate:
+                    sleep_for = self.heartrate - time_since_last_heartbeat
+                    self.log.warning("Time since last heartbeat(%.2f s) < heartrate(%s s)"
+                                     ", sleeping for %s s", time_since_last_heartbeat,
+                                     self.heartrate, sleep_for)
+                    time.sleep(sleep_for)
         finally:
             self.on_kill()
 

airflow/utils/helpers.py

@@ -351,3 +351,7 @@ def render_log_filename(ti, try_number, filename_template):
                                     task_id=ti.task_id,
                                     execution_date=ti.execution_date.isoformat(),
                                     try_number=try_number)
+
+
+def convert_camel_to_snake(camel_str):
+    return re.sub('(?!^)([A-Z]+)', r'_\1', camel_str).lower()

tests/dags/test_heartbeat_failed_fast.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 datetime import datetime
+
+from airflow.models import DAG
+from airflow.operators.bash_operator import BashOperator
+
+DEFAULT_DATE = datetime(2016, 1, 1)
+
+args = {
+    'owner': 'airflow',
+    'start_date': DEFAULT_DATE,
+}
+
+dag = DAG(dag_id='test_heartbeat_failed_fast', default_args=args)
+task = BashOperator(
+    task_id='test_heartbeat_failed_fast_op',
+    bash_command='sleep 5',
+    dag=dag)

tests/jobs/test_local_task_job.py

@@ -30,10 +30,12 @@
 from airflow.operators.dummy_operator import DummyOperator
 from airflow.task.task_runner.base_task_runner import BaseTaskRunner
 from airflow.utils import timezone
+from airflow.utils.db import create_session
 from airflow.utils.net import get_hostname
 from airflow.utils.state import State
 from tests.compat import patch
 from tests.core import TEST_DAG_FOLDER
+from tests.executors.test_executor import TestExecutor
 from tests.test_utils.db import clear_db_runs
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
@@ -116,6 +118,50 @@ def test_localtaskjob_heartbeat(self, mock_pid):
         mock_pid.return_value = 2
         self.assertRaises(AirflowException, job1.heartbeat_callback)
 
+    @patch('os.getpid')
+    def test_heartbeat_failed_fast(self, mock_getpid):
+        """
+        Test that task heartbeat will sleep when it fails fast
+        """
+        mock_getpid.return_value = 1
+
+        heartbeat_records = []
+
+        def heartbeat_recorder():
+            heartbeat_records.append(timezone.utcnow())
+
+        with create_session() as session:
+            dagbag = models.DagBag(
+                dag_folder=TEST_DAG_FOLDER,
+                include_examples=False,
+            )
+            dag_id = 'test_heartbeat_failed_fast'
+            task_id = 'test_heartbeat_failed_fast_op'
+            dag = dagbag.get_dag(dag_id)
+            task = dag.get_task(task_id)
+
+            dag.create_dagrun(run_id="test_heartbeat_failed_fast_run",
+                              state=State.RUNNING,
+                              execution_date=DEFAULT_DATE,
+                              start_date=DEFAULT_DATE,
+                              session=session)
+            ti = TI(task=task, execution_date=DEFAULT_DATE)
+            ti.refresh_from_db()
+            ti.state = State.RUNNING
+            ti.hostname = get_hostname()
+            ti.pid = 1
+            session.commit()
+
+            job = LocalTaskJob(task_instance=ti, executor=TestExecutor(do_update=False))
+            job.heartrate = 2
+            job.heartbeat = heartbeat_recorder
+            job._execute()
+            self.assertGreater(len(heartbeat_records), 1)
+            for i in range(1, len(heartbeat_records)):
+                time1 = heartbeat_records[i - 1]
+                time2 = heartbeat_records[i]
+                self.assertGreaterEqual((time2 - time1).total_seconds(), job.heartrate)
+
     @unittest.skipIf('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
                      "flaky when run on mysql")
     @unittest.skipIf('postgresql' in configuration.conf.get('core', 'sql_alchemy_conn'),

tests/utils/test_helpers.py

@@ -270,6 +270,11 @@ def test_chain_different_length_iterable(self):
         with self.assertRaises(AirflowException):
             helpers.chain([t1, t2], [t3, t4, t5])
 
+    def test_convert_camel_to_snake(self):
+        self.assertEqual(helpers.convert_camel_to_snake('LocalTaskJob'), 'local_task_job')
+        self.assertEqual(helpers.convert_camel_to_snake('somethingVeryRandom'),
+                         'something_very_random')
+
 
 if __name__ == '__main__':
     unittest.main()