Repository: apache/airflow

Issue:
AirflowTaskTimeout should inherit BaseException instead of Exception?
### Discussed in https://github.com/apache/airflow/discussions/35643

<div type='discussions-op-text'>

<sup>Originally posted by **hterik** November 14, 2023</sup>
### Apache Airflow version

2.7.3

### What happened

Assume following code runs inside airflow task:

```py
errors = []
for item in items_to_process:
    try: 
        process(item)   # This can take several minutes/hours each
    except Exception as e:
        errors.append(e)
if errors:
    raise Exception(......)
```

If the Airflow task time out, what happens is that it injects an `AirflowTaskTimeout` exception where the code is currently running now.
If the code is designed to capture exceptions, it will capture the timeout and potentially continue running for several hours anyway.


### What you think should happen instead

I think it would be better if `AirflowTaskTimeout` was treated similarly to `KeyboardInterrupt`, so it need explicit `except` block to capture. 
Moving it outside of the `Exception` inheritance tree to inherit directly from `BaseException` would solve that.
See https://docs.python.org/3/library/exceptions.html#exception-hierarchy 
Timeouts are generally in place to be last resort of aborting process by the higher level system if the lower level code has bugs in it, so it's safer to assume that code that doesn't deal with catching timeout should be aborted.

### How to reproduce

See above

### Operating System

NA

### Versions of Apache Airflow Providers

apache-airflow==2.7.1


### Deployment

Official Apache Airflow Helm Chart

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
</div>


Discussion:
Yes. good idea. Might be worth implementing it. 
BTW. There is another problem with tasks running low-level C-code that are not interrupted by signal (and should be handled in yet another way) - see  https://github.com/apache/airflow/issues/35474  - but indeed when code is catching all exceptions, chaning timeout to BaseException should work well. Would you like to implement it @hterik ?
That is exactly what a **Option 1** propose 😄 
- https://github.com/apache/airflow/issues/35474#issuecomment-1801592051

Resolving patch:
diff --git a/airflow/exceptions.py b/airflow/exceptions.py
--- a/airflow/exceptions.py
+++ b/airflow/exceptions.py
@@ -79,7 +79,10 @@ class InvalidStatsNameException(AirflowException):
     """Raise when name of the stats is invalid."""
 
 
-class AirflowTaskTimeout(AirflowException):
+# Important to inherit BaseException instead of AirflowException->Exception, since this Exception is used
+# to explicitly interrupt ongoing task. Code that does normal error-handling should not treat
+# such interrupt as an error that can be handled normally. (Compare with KeyboardInterrupt)
+class AirflowTaskTimeout(BaseException):
     """Raise when the task execution times-out."""
 
 
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -812,7 +812,7 @@ def _is_eligible_to_retry(*, task_instance: TaskInstance | TaskInstancePydantic)
 def _handle_failure(
     *,
     task_instance: TaskInstance | TaskInstancePydantic,
-    error: None | str | Exception | KeyboardInterrupt,
+    error: None | str | BaseException,
     session: Session,
     test_mode: bool | None = None,
     context: Context | None = None,
@@ -2411,7 +2411,7 @@ def _run_raw_task(
                 self.handle_failure(e, test_mode, context, force_fail=True, session=session)
                 session.commit()
                 raise
-            except AirflowException as e:
+            except (AirflowTaskTimeout, AirflowException) as e:
                 if not test_mode:
                     self.refresh_from_db(lock_for_update=True, session=session)
                 # for case when task is marked as success/failed externally
@@ -2426,10 +2426,6 @@ def _run_raw_task(
                     self.handle_failure(e, test_mode, context, session=session)
                     session.commit()
                     raise
-            except (Exception, KeyboardInterrupt) as e:
-                self.handle_failure(e, test_mode, context, session=session)
-                session.commit()
-                raise
             except SystemExit as e:
                 # We have already handled SystemExit with success codes (0 and None) in the `_execute_task`.
                 # Therefore, here we must handle only error codes.
@@ -2437,6 +2433,10 @@ def _run_raw_task(
                 self.handle_failure(msg, test_mode, context, session=session)
                 session.commit()
                 raise Exception(msg)
+            except BaseException as e:
+                self.handle_failure(e, test_mode, context, session=session)
+                session.commit()
+                raise
             finally:
                 Stats.incr(f"ti.finish.{self.dag_id}.{self.task_id}.{self.state}", tags=self.stats_tags)
                 # Same metric with tagging
@@ -2743,7 +2743,7 @@ def get_truncated_error_traceback(error: BaseException, truncate_to: Callable) -
     def fetch_handle_failure_context(
         cls,
         ti: TaskInstance | TaskInstancePydantic,
-        error: None | str | Exception | KeyboardInterrupt,
+        error: None | str | BaseException,
         test_mode: bool | None = None,
         context: Context | None = None,
         force_fail: bool = False,
@@ -2838,7 +2838,7 @@ def save_to_db(ti: TaskInstance | TaskInstancePydantic, session: Session = NEW_S
     @provide_session
     def handle_failure(
         self,
-        error: None | str | Exception | KeyboardInterrupt,
+        error: None | str | BaseException,
         test_mode: bool | None = None,
         context: Context | None = None,
         force_fail: bool = False,
diff --git a/airflow/providers/celery/executors/celery_executor_utils.py b/airflow/providers/celery/executors/celery_executor_utils.py
--- a/airflow/providers/celery/executors/celery_executor_utils.py
+++ b/airflow/providers/celery/executors/celery_executor_utils.py
@@ -41,7 +41,7 @@
 
 import airflow.settings as settings
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowTaskTimeout
 from airflow.executors.base_executor import BaseExecutor
 from airflow.stats import Stats
 from airflow.utils.dag_parsing_context import _airflow_parsing_context_manager
@@ -198,7 +198,7 @@ class ExceptionWithTraceback:
     :param exception_traceback: The stacktrace to wrap
     """
 
-    def __init__(self, exception: Exception, exception_traceback: str):
+    def __init__(self, exception: BaseException, exception_traceback: str):
         self.exception = exception
         self.traceback = exception_traceback
 
@@ -211,7 +211,7 @@ def send_task_to_executor(
     try:
         with timeout(seconds=OPERATION_TIMEOUT):
             result = task_to_run.apply_async(args=[command], queue=queue)
-    except Exception as e:
+    except (Exception, AirflowTaskTimeout) as e:
         exception_traceback = f"Celery Task ID: {key}\n{traceback.format_exc()}"
         result = ExceptionWithTraceback(e, exception_traceback)
 


Test patch:
diff --git a/tests/core/test_core.py b/tests/core/test_core.py
--- a/tests/core/test_core.py
+++ b/tests/core/test_core.py
@@ -71,11 +71,18 @@ class InvalidTemplateFieldOperator(BaseOperator):
             op.dry_run()
 
     def test_timeout(self, dag_maker):
+        def sleep_and_catch_other_exceptions():
+            try:
+                sleep(5)
+                # Catching Exception should NOT catch AirflowTaskTimeout
+            except Exception:
+                pass
+
         with dag_maker():
             op = PythonOperator(
                 task_id="test_timeout",
                 execution_timeout=timedelta(seconds=1),
-                python_callable=lambda: sleep(5),
+                python_callable=sleep_and_catch_other_exceptions,
             )
         dag_maker.create_dagrun()
         with pytest.raises(AirflowTaskTimeout):
diff --git a/tests/providers/microsoft/azure/hooks/test_synapse.py b/tests/providers/microsoft/azure/hooks/test_synapse.py
--- a/tests/providers/microsoft/azure/hooks/test_synapse.py
+++ b/tests/providers/microsoft/azure/hooks/test_synapse.py
@@ -21,6 +21,7 @@
 import pytest
 from azure.synapse.spark import SparkClient
 
+from airflow.exceptions import AirflowTaskTimeout
 from airflow.models.connection import Connection
 from airflow.providers.microsoft.azure.hooks.synapse import AzureSynapseHook, AzureSynapseSparkBatchRunStatus
 
@@ -172,7 +173,7 @@ def test_wait_for_job_run_status(hook, job_run_status, expected_status, expected
         if expected_output != "timeout":
             assert hook.wait_for_job_run_status(**config) == expected_output
         else:
-            with pytest.raises(Exception):
+            with pytest.raises(AirflowTaskTimeout):
                 hook.wait_for_job_run_status(**config)