From 77cb73dca95df40d827fc56864fd380d248ab37d Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Fri, 12 Jun 2020 16:27:39 +0100 Subject: [PATCH] Mark task as failed when it fails sending in Celery If a task failed hard on celery, _before_ being able to execute the airflow code the task would end up stuck in queued state. This change makes it get retried. This was discovered in load testing the HA work (but unrelated to HA changes), where I swamped the kube-dns pod, meaning the worker was sometimes unable to resolve the db name via DNS, so the state in the DB was never updated --- airflow/executors/celery_executor.py | 6 +++--- tests/executors/test_celery_executor.py | 7 ++++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/airflow/executors/celery_executor.py b/airflow/executors/celery_executor.py index 600183697c013..13726674bd2b0 100644 --- a/airflow/executors/celery_executor.py +++ b/airflow/executors/celery_executor.py @@ -42,6 +42,7 @@ from airflow.models.taskinstance import SimpleTaskInstance, TaskInstanceKey from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname +from airflow.utils.state import State from airflow.utils.timeout import timeout log = logging.getLogger(__name__) @@ -188,14 +189,13 @@ def _process_tasks(self, task_tuples_to_send: List[TaskInstanceInCelery]) -> Non self.log.debug('Sent all tasks.') for key, _, result in key_and_async_results: + self.queued_tasks.pop(key) if isinstance(result, ExceptionWithTraceback): self.log.error( # pylint: disable=logging-not-lazy CELERY_SEND_ERR_MSG_HEADER + ":%s\n%s\n", result.exception, result.traceback ) + self.event_buffer[key] = (State.FAILED, None) elif result is not None: - # Only pops when enqueued successfully, otherwise keep it - # and expect scheduler loop to deal with it. - self.queued_tasks.pop(key) result.backend = cached_celery_backend self.running.add(key) self.tasks[key] = result diff --git a/tests/executors/test_celery_executor.py b/tests/executors/test_celery_executor.py index 76c2ace58e734..76b572e260eab 100644 --- a/tests/executors/test_celery_executor.py +++ b/tests/executors/test_celery_executor.py @@ -159,13 +159,14 @@ def fake_execute_command(): dag=DAG(dag_id='id'), start_date=datetime.datetime.now() ) + when = datetime.datetime.now() value_tuple = 'command', 1, None, \ SimpleTaskInstance(ti=TaskInstance(task=task, execution_date=datetime.datetime.now())) - key = ('fail', 'fake_simple_ti', datetime.datetime.now(), 0) + key = ('fail', 'fake_simple_ti', when, 0) executor.queued_tasks[key] = value_tuple executor.heartbeat() - self.assertEqual(1, len(executor.queued_tasks)) - self.assertEqual(executor.queued_tasks[key], value_tuple) + self.assertEqual(0, len(executor.queued_tasks), "Task should no longer be queued") + self.assertEqual(executor.event_buffer[('fail', 'fake_simple_ti', when, 0)][0], State.FAILED) @pytest.mark.backend("mysql", "postgres") def test_exception_propagation(self):