Airflow exceptions thrown on on_retry_callback is suppressed - airflow

Question

We are using on_retry_callback parameter available in the Airflow operators to do some cleanup activities before the task is retried. If there are exceptions thrown on the on_retry_callback function, the exceptions are not logged in the task_instance's log. Without the exception details, it is getting difficult to debug if there are issues in the on_retry_callback function. If this is the default behavior, is there a workaround to enable logging for the exceptions?.
Note: We are using the airflow 2.0.2 version.
Please let me know if there are any questions.
Sample Dag to explain this is given below.
from datetime import datetime
from airflow.operators.python import PythonOperator
from airflow.models.dag import DAG
def sample_function2():
var = 1 / 0
def on_retry_callback_sample(context):
print(f'on_retry_callback_started')
v = 1 / 0
print(f'on_retry_callback completed')
dag = DAG(
'venkat-test-dag',
description='This is a test dag',
start_date=datetime(2023, 1, 10, 18, 0),
schedule_interval='0 12 * * *',
catchup=False
)
func2 = PythonOperator(task_id='function2',
python_callable=sample_function2,
dag=dag,
retries=2,
on_retry_callback=on_retry_callback_sample)
func2
Log file of this run on the local airflow setup is given below. If you see the last message we see on the log file "on_retry_callback_started" but I expect some ZeroDivisionError after this line and finally the line "on_retry_callback completed". How can I achieve this?.
14f0fed99882
*** Reading local file: /usr/local/airflow/logs/venkat-test-dag/function2/2023-01-13T13:22:03.178261+00:00/1.log
[2023-01-13 13:22:05,091] {{taskinstance.py:877}} INFO - Dependencies all met for <TaskInstance: venkat-test-dag.function2 2023-01-13T13:22:03.178261+00:00 [queued]>
[2023-01-13 13:22:05,128] {{taskinstance.py:877}} INFO - Dependencies all met for <TaskInstance: venkat-test-dag.function2 2023-01-13T13:22:03.178261+00:00 [queued]>
[2023-01-13 13:22:05,128] {{taskinstance.py:1068}} INFO -
--------------------------------------------------------------------------------
[2023-01-13 13:22:05,128] {{taskinstance.py:1069}} INFO - Starting attempt 1 of 3
[2023-01-13 13:22:05,128] {{taskinstance.py:1070}} INFO -
--------------------------------------------------------------------------------
[2023-01-13 13:22:05,143] {{taskinstance.py:1089}} INFO - Executing <Task(PythonOperator): function2> on 2023-01-13T13:22:03.178261+00:00
[2023-01-13 13:22:05,145] {{standard_task_runner.py:52}} INFO - Started process 6947 to run task
[2023-01-13 13:22:05,150] {{standard_task_runner.py:76}} INFO - Running: ['airflow', 'tasks', 'run', 'venkat-test-dag', 'function2', '2023-01-13T13:22:03.178261+00:00', '--job-id', '356', '--pool', 'default_pool', '--raw', '--subdir', 'DAGS_FOLDER/dp-etl-mixpanel_stg-24H/dags/venkat-test-dag.py', '--cfg-path', '/tmp/tmpny0mhh4j', '--error-file', '/tmp/tmpul506kro']
[2023-01-13 13:22:05,151] {{standard_task_runner.py:77}} INFO - Job 356: Subtask function2
[2023-01-13 13:22:05,244] {{logging_mixin.py:104}} INFO - Running <TaskInstance: venkat-test-dag.function2 2023-01-13T13:22:03.178261+00:00 [running]> on host 14f0fed99882
[2023-01-13 13:22:05,345] {{taskinstance.py:1283}} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=venkat-test-dag
AIRFLOW_CTX_TASK_ID=function2
AIRFLOW_CTX_EXECUTION_DATE=2023-01-13T13:22:03.178261+00:00
AIRFLOW_CTX_DAG_RUN_ID=manual__2023-01-13T13:22:03.178261+00:00
[2023-01-13 13:22:05,346] {{taskinstance.py:1482}} ERROR - Task failed with exception
Traceback (most recent call last):
File "/usr/local/lib/python3.7/site-packages/airflow/models/taskinstance.py", line 1138, in _run_raw_task
self._prepare_and_execute_task_with_callbacks(context, task)
File "/usr/local/lib/python3.7/site-packages/airflow/models/taskinstance.py", line 1311, in _prepare_and_execute_task_with_callbacks
result = self._execute_task(context, task_copy)
File "/usr/local/lib/python3.7/site-packages/airflow/models/taskinstance.py", line 1341, in _execute_task
result = task_copy.execute(context=context)
File "/usr/local/lib/python3.7/site-packages/airflow/operators/python.py", line 117, in execute
return_value = self.execute_callable()
File "/usr/local/lib/python3.7/site-packages/airflow/operators/python.py", line 128, in execute_callable
return self.python_callable(*self.op_args, **self.op_kwargs)
File "/usr/local/airflow/dags/dp-etl-mixpanel_stg-24H/dags/venkat-test-dag.py", line 7, in sample_function2
var = 1 / 0
ZeroDivisionError: division by zero
[2023-01-13 13:22:05,349] {{taskinstance.py:1532}} INFO - Marking task as UP_FOR_RETRY. dag_id=venkat-test-dag, task_id=function2, execution_date=20230113T132203, start_date=20230113T132205, end_date=20230113T132205
[2023-01-13 13:22:05,402] {{local_task_job.py:146}} INFO - Task exited with return code 1
[2023-01-13 13:22:05,459] {{logging_mixin.py:104}} INFO - on_retry_callback_started

Answer 1

Adding as an answer for visibility:
This issue is likely related to a fix which was merged in Airflow version 2.1.3:
https://github.com/apache/airflow/pull/17347