From 91076843d6e3de9326acc9207495a106a4c41968 Mon Sep 17 00:00:00 2001 From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com> Date: Fri, 22 Oct 2021 23:21:53 -0400 Subject: [PATCH 1/4] Added ti.clear_next_method_args() to clear ti.next_method and ti.next_kwargs when a task finishes running. --- airflow/models/taskinstance.py | 18 ++++- tests/models/test_taskinstance.py | 125 ++++++++++++++++++++++++++++-- 2 files changed, 134 insertions(+), 9 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 6289ecaf06dd5..f9690371757c6 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -1278,6 +1278,14 @@ def _log_state(self, lead_msg: str = ''): self._date_or_empty('end_date'), ) + # Ensure we unset next_method and next_kwargs to ensure that any + # retries don't re-use them. + def clear_next_method_args(self): + self.log.debug("Clearing next_method and next_kwargs.") + + self.next_method = None + self.next_kwargs = None + @provide_session @Sentry.enrich_errors def _run_raw_task( @@ -1367,6 +1375,8 @@ def _run_raw_task( # or dagrun timed out and task is marked as skipped # current behavior doesn't hit the callbacks if self.state in State.finished: + self.clear_next_method_args() + session.commit() return else: self.handle_failure(e, test_mode, error_file=error_file, session=session) @@ -1380,6 +1390,7 @@ def _run_raw_task( Stats.incr(f'ti.finish.{self.task.dag_id}.{self.task.task_id}.{self.state}') # Recording SKIPPED or SUCCESS + self.clear_next_method_args() self.end_date = timezone.utcnow() self._log_state() self.set_duration() @@ -1665,6 +1676,8 @@ def _handle_reschedule(self, actual_start_date, reschedule_exception, test_mode= # to same log file. self._try_number -= 1 + self.clear_next_method_args() + session.merge(self) session.commit() self.log.info('Rescheduling task, marking task as UP_FOR_RESCHEDULE') @@ -1706,10 +1719,7 @@ def handle_failure( dag_run = self.get_dagrun(session=session) # self.dag_run not populated by refresh_from_db session.add(TaskFail(task, dag_run.execution_date, self.start_date, self.end_date)) - # Ensure we unset next_method and next_kwargs to ensure that any - # retries don't re-use them. - self.next_method = None - self.next_kwargs = None + self.clear_next_method_args() # Set state correctly and figure out how to log it and decide whether # to email diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index d22aef8c1c7b1..af60770eb4c55 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -33,6 +33,7 @@ from airflow.exceptions import ( AirflowException, AirflowFailException, + AirflowRescheduleException, AirflowSensorTimeout, AirflowSkipException, ) @@ -481,16 +482,71 @@ def task_function(ti): ti.refresh_from_db() ti.state == state - def test_task_retry_wipes_next_fields(self, session, dag_maker): + def test_task_wipes_next_fields_success(self, session, dag_maker): """ Test that ensures that tasks wipe their next_method and next_kwargs - fields when they are queued for retry after a failure. + when they go into a state of SUCCESS. """ - with dag_maker('test_mark_failure_2'): + with dag_maker("test_deferred_method_clear"): task = BashOperator( - task_id='test_retry_handling_op', - bash_command='exit 1', + task_id="test_deferred_method_clear_task", + bash_command="" + ) + + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + ti.next_method = "execute" + ti.next_kwargs = {} + session.merge(ti) + session.commit() + + ti.task = task + ti.run() + ti.refresh_from_db() + + assert ti.next_method is None + assert ti.next_kwargs is None + assert ti.state == State.SUCCESS + + def test_task_wipes_next_fields_failure(self, session, dag_maker): + """ + Test that ensures that tasks wipe their next_method and next_kwargs + when they go into a state of FAILED. + """ + + with dag_maker("test_deferred_method_clear"): + task = BashOperator( + task_id="test_deferred_method_clear_task", + bash_command="exit 1" + ) + + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + ti.next_method = "execute" + ti.next_kwargs = {} + session.merge(ti) + session.commit() + + ti.task = task + with pytest.raises(AirflowException): + ti.run() + ti.refresh_from_db() + + assert ti.next_method is None + assert ti.next_kwargs is None + assert ti.state == State.FAILED + + def test_task_wipes_next_fields_up_for_retry(self, session, dag_maker): + """ + Test that ensures that tasks wipe their next_method and next_kwargs + when they go into a state of UP_FOR_RETRY. + """ + + with dag_maker("test_deferred_method_clear"): + task = BashOperator( + task_id="test_deferred_method_clear_task", + bash_command="exit 1", retries=1, retry_delay=datetime.timedelta(seconds=2), ) @@ -511,6 +567,65 @@ def test_task_retry_wipes_next_fields(self, session, dag_maker): assert ti.next_kwargs is None assert ti.state == State.UP_FOR_RETRY + def test_task_wipes_next_fields_skipped(self, session, dag_maker): + """ + Test that ensures that tasks wipe their next_method and next_kwargs + when they go into a state of SKIPPED. + """ + def skip(): + raise AirflowSkipException + + with dag_maker("test_deferred_method_clear"): + task = PythonOperator( + task_id="test_deferred_method_clear_task", + python_callable=skip, + ) + + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + ti.next_method = "execute" + ti.next_kwargs = {} + session.merge(ti) + session.commit() + + ti.task = task + ti.run() + ti.refresh_from_db() + + assert ti.next_method is None + assert ti.next_kwargs is None + assert ti.state == State.SKIPPED + + def test_task_wipes_next_fields_reschedule(self, session, dag_maker): + """ + Test that ensures that tasks wipe their next_method and next_kwargs + when they go into a state of SKIPPED. + """ + def reschedule(): + reschedule_date = timezone.utcnow() + raise AirflowRescheduleException(reschedule_date) + + with dag_maker("test_deferred_method_clear"): + task = PythonOperator( + task_id="test_deferred_method_clear_task", + python_callable=reschedule, + ) + + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + ti.next_method = "execute" + ti.next_kwargs = {} + session.merge(ti) + session.commit() + + ti.task = task + ti.run() + ti.refresh_from_db() + + assert ti.next_method is None + assert ti.next_kwargs is None + assert ti.state == State.UP_FOR_RESCHEDULE + @freeze_time('2021-09-19 04:56:35', as_kwarg='frozen_time') def test_retry_delay(self, dag_maker, frozen_time=None): """ From 2a3bbe79c414f447f242c9a9766a2bd840d0487b Mon Sep 17 00:00:00 2001 From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com> Date: Sat, 23 Oct 2021 10:25:40 -0400 Subject: [PATCH 2/4] Updated formatting. --- tests/models/test_taskinstance.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index af60770eb4c55..20ac1a92571c8 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -489,10 +489,7 @@ def test_task_wipes_next_fields_success(self, session, dag_maker): """ with dag_maker("test_deferred_method_clear"): - task = BashOperator( - task_id="test_deferred_method_clear_task", - bash_command="" - ) + task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="") dr = dag_maker.create_dagrun() ti = dr.task_instances[0] @@ -516,10 +513,7 @@ def test_task_wipes_next_fields_failure(self, session, dag_maker): """ with dag_maker("test_deferred_method_clear"): - task = BashOperator( - task_id="test_deferred_method_clear_task", - bash_command="exit 1" - ) + task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="exit 1") dr = dag_maker.create_dagrun() ti = dr.task_instances[0] @@ -572,6 +566,7 @@ def test_task_wipes_next_fields_skipped(self, session, dag_maker): Test that ensures that tasks wipe their next_method and next_kwargs when they go into a state of SKIPPED. """ + def skip(): raise AirflowSkipException @@ -601,6 +596,7 @@ def test_task_wipes_next_fields_reschedule(self, session, dag_maker): Test that ensures that tasks wipe their next_method and next_kwargs when they go into a state of SKIPPED. """ + def reschedule(): reschedule_date = timezone.utcnow() raise AirflowRescheduleException(reschedule_date) From e338f6cdd297c41b79b479828ef641b152a7b3d9 Mon Sep 17 00:00:00 2001 From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com> Date: Sat, 23 Oct 2021 11:20:34 -0400 Subject: [PATCH 3/4] Consolidated into a Parameterized test. --- tests/models/test_taskinstance.py | 144 ++++++++---------------------- 1 file changed, 35 insertions(+), 109 deletions(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 20ac1a92571c8..211e67f662417 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -482,129 +482,51 @@ def task_function(ti): ti.refresh_from_db() ti.state == state - def test_task_wipes_next_fields_success(self, session, dag_maker): - """ - Test that ensures that tasks wipe their next_method and next_kwargs - when they go into a state of SUCCESS. - """ - - with dag_maker("test_deferred_method_clear"): - task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="") - - dr = dag_maker.create_dagrun() - ti = dr.task_instances[0] - ti.next_method = "execute" - ti.next_kwargs = {} - session.merge(ti) - session.commit() - - ti.task = task - ti.run() - ti.refresh_from_db() - - assert ti.next_method is None - assert ti.next_kwargs is None - assert ti.state == State.SUCCESS - - def test_task_wipes_next_fields_failure(self, session, dag_maker): + @pytest.mark.parametrize( + "state", + [State.FAILED, State.SKIPPED, State.SUCCESS, State.UP_FOR_RESCHEDULE, State.UP_FOR_RETRY], + ) + def test_task_wipes_next_fields(self, session, state, dag_maker): """ Test that ensures that tasks wipe their next_method and next_kwargs - when they go into a state of FAILED. + when they go into a state of FAILED, SKIPPED, SUCCESS, UP_FOR_RESCHEDULE, or UP_FOR_RETRY. """ - with dag_maker("test_deferred_method_clear"): - task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="exit 1") - - dr = dag_maker.create_dagrun() - ti = dr.task_instances[0] - ti.next_method = "execute" - ti.next_kwargs = {} - session.merge(ti) - session.commit() - - ti.task = task - with pytest.raises(AirflowException): - ti.run() - ti.refresh_from_db() - - assert ti.next_method is None - assert ti.next_kwargs is None - assert ti.state == State.FAILED - - def test_task_wipes_next_fields_up_for_retry(self, session, dag_maker): - """ - Test that ensures that tasks wipe their next_method and next_kwargs - when they go into a state of UP_FOR_RETRY. - """ - - with dag_maker("test_deferred_method_clear"): - task = BashOperator( - task_id="test_deferred_method_clear_task", - bash_command="exit 1", - retries=1, - retry_delay=datetime.timedelta(seconds=2), - ) - - dr = dag_maker.create_dagrun() - ti = dr.task_instances[0] - ti.next_method = "execute" - ti.next_kwargs = {} - session.merge(ti) - session.commit() - - ti.task = task - with pytest.raises(AirflowException): - ti.run() - ti.refresh_from_db() - - assert ti.next_method is None - assert ti.next_kwargs is None - assert ti.state == State.UP_FOR_RETRY - - def test_task_wipes_next_fields_skipped(self, session, dag_maker): - """ - Test that ensures that tasks wipe their next_method and next_kwargs - when they go into a state of SKIPPED. - """ + def failure(): + raise AirflowException def skip(): raise AirflowSkipException - with dag_maker("test_deferred_method_clear"): - task = PythonOperator( - task_id="test_deferred_method_clear_task", - python_callable=skip, - ) - - dr = dag_maker.create_dagrun() - ti = dr.task_instances[0] - ti.next_method = "execute" - ti.next_kwargs = {} - session.merge(ti) - session.commit() - - ti.task = task - ti.run() - ti.refresh_from_db() - - assert ti.next_method is None - assert ti.next_kwargs is None - assert ti.state == State.SKIPPED - - def test_task_wipes_next_fields_reschedule(self, session, dag_maker): - """ - Test that ensures that tasks wipe their next_method and next_kwargs - when they go into a state of SKIPPED. - """ + def success(): + return None def reschedule(): reschedule_date = timezone.utcnow() raise AirflowRescheduleException(reschedule_date) + _retries = 0 + _retry_delay = datetime.timedelta(seconds=0) + + if state == State.FAILED: + _python_callable = failure + elif state == State.SKIPPED: + _python_callable = skip + elif state == State.SUCCESS: + _python_callable = success + elif state == State.UP_FOR_RESCHEDULE: + _python_callable = reschedule + elif state in [State.FAILED, State.UP_FOR_RETRY]: + _python_callable = failure + _retries = 1 + _retry_delay = datetime.timedelta(seconds=2) + with dag_maker("test_deferred_method_clear"): task = PythonOperator( task_id="test_deferred_method_clear_task", - python_callable=reschedule, + python_callable=_python_callable, + retries=_retries, + retry_delay=_retry_delay, ) dr = dag_maker.create_dagrun() @@ -615,12 +537,16 @@ def reschedule(): session.commit() ti.task = task - ti.run() + if state in [State.FAILED, State.UP_FOR_RETRY]: + with pytest.raises(AirflowException): + ti.run() + elif state in [State.SKIPPED, State.SUCCESS, State.UP_FOR_RESCHEDULE]: + ti.run() ti.refresh_from_db() assert ti.next_method is None assert ti.next_kwargs is None - assert ti.state == State.UP_FOR_RESCHEDULE + assert ti.state == state @freeze_time('2021-09-19 04:56:35', as_kwarg='frozen_time') def test_retry_delay(self, dag_maker, frozen_time=None): From 296a350abf0b5b531d218564fd718b0cb745d625 Mon Sep 17 00:00:00 2001 From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com> Date: Sat, 23 Oct 2021 11:24:45 -0400 Subject: [PATCH 4/4] Added session.merge to match other cases. --- airflow/models/taskinstance.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index f9690371757c6..38ae39727eb01 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -1376,6 +1376,7 @@ def _run_raw_task( # current behavior doesn't hit the callbacks if self.state in State.finished: self.clear_next_method_args() + session.merge(self) session.commit() return else: