AIP-51 - Misc. Compatibility Checks

airflow/executors/base_executor.py

@@ -106,6 +106,8 @@ class BaseExecutor(LoggingMixin):
     """
 
     supports_ad_hoc_ti_run: bool = False
+    supports_pickling: bool = True
+    supports_sentry: bool = False
 
     job_id: None | int | str = None
     callback_sink: BaseCallbackSink | None = None

airflow/executors/celery_executor.py

@@ -228,6 +228,7 @@ class CeleryExecutor(BaseExecutor):
     """
 
     supports_ad_hoc_ti_run: bool = True
+    supports_sentry: bool = True
 
     def __init__(self):
         super().__init__()

airflow/executors/celery_kubernetes_executor.py

@@ -39,6 +39,9 @@ class CeleryKubernetesExecutor(LoggingMixin):
     """
 
     supports_ad_hoc_ti_run: bool = True
+    supports_pickling: bool = True
+    supports_sentry: bool = False
+
     callback_sink: BaseCallbackSink | None = None
 
     KUBERNETES_QUEUE = conf.get("celery_kubernetes_executor", "kubernetes_queue")

airflow/executors/dask_executor.py

@@ -43,6 +43,8 @@
 class DaskExecutor(BaseExecutor):
     """DaskExecutor submits tasks to a Dask Distributed cluster."""
 
+    supports_pickling: bool = False
+
     def __init__(self, cluster_address=None):
         super().__init__(parallelism=0)
         if cluster_address is None:

airflow/executors/executor_loader.py

@@ -169,6 +169,8 @@ def __load_local_kubernetes_executor(cls) -> BaseExecutor:
         return local_kubernetes_executor_cls(local_executor, kubernetes_executor)
 
 
+# This tuple is deprecated due to AIP-51 and is no longer used in core Airflow.
+# TODO: Remove in Airflow 3.0
 UNPICKLEABLE_EXECUTORS = (
     LOCAL_EXECUTOR,
     SEQUENTIAL_EXECUTOR,

airflow/executors/local_executor.py

@@ -206,6 +206,7 @@ class LocalExecutor(BaseExecutor):
     """
 
     is_local: bool = True
+    supports_pickling: bool = False
 
     def __init__(self, parallelism: int = PARALLELISM):
         super().__init__(parallelism=parallelism)

airflow/executors/local_kubernetes_executor.py

@@ -39,6 +39,9 @@ class LocalKubernetesExecutor(LoggingMixin):
     """
 
     supports_ad_hoc_ti_run: bool = True
+    supports_pickling: bool = False
+    supports_sentry: bool = False
+
     callback_sink: BaseCallbackSink | None = None
 
     KUBERNETES_QUEUE = conf.get("local_kubernetes_executor", "kubernetes_queue")

airflow/executors/sequential_executor.py

@@ -44,6 +44,7 @@ class SequentialExecutor(BaseExecutor):
     SequentialExecutor alongside sqlite as you first install it.
     """
 
+    supports_pickling: bool = False
     is_local: bool = True
 
     def __init__(self):

airflow/jobs/backfill_job.py

@@ -35,7 +35,6 @@
     PoolNotFound,
     TaskConcurrencyLimitReached,
 )
-from airflow.executors import executor_constants
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.jobs.base_job import BaseJob
 from airflow.models import DAG, DagPickle
@@ -845,11 +844,9 @@ def _execute(self, session=None):
         # picklin'
         pickle_id = None
 
-        if not self.donot_pickle and self.executor_class not in (
-            executor_constants.LOCAL_EXECUTOR,
-            executor_constants.SEQUENTIAL_EXECUTOR,
-            executor_constants.DASK_EXECUTOR,
-        ):
+        executor_class, _ = ExecutorLoader.import_default_executor_cls()
+
+        if not self.donot_pickle and executor_class.supports_pickling:
             pickle = DagPickle(self.dag)
             session.add(pickle)
             session.commit()

airflow/jobs/scheduler_job.py

@@ -41,7 +41,7 @@
 from airflow.callbacks.pipe_callback_sink import PipeCallbackSink
 from airflow.configuration import conf
 from airflow.exceptions import RemovedInAirflow3Warning
-from airflow.executors.executor_loader import UNPICKLEABLE_EXECUTORS
+from airflow.executors.executor_loader import ExecutorLoader
 from airflow.jobs.base_job import BaseJob
 from airflow.models.dag import DAG, DagModel
 from airflow.models.dagbag import DagBag
@@ -720,8 +720,10 @@ def _execute(self) -> None:
 
         self.log.info("Starting the scheduler")
 
+        executor_class, _ = ExecutorLoader.import_default_executor_cls()
+
         # DAGs can be pickled for easier remote execution by some executors
-        pickle_dags = self.do_pickle and self.executor_class not in UNPICKLEABLE_EXECUTORS
+        pickle_dags = self.do_pickle and executor_class.supports_pickling
 
         self.log.info("Processing each file at most %s times", self.num_times_parse_dags)
 

airflow/sentry.py

@@ -23,6 +23,7 @@
 from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
+from airflow.executors.executor_loader import ExecutorLoader
 from airflow.utils.session import find_session_idx, provide_session
 from airflow.utils.state import State
 
@@ -82,14 +83,15 @@ class ConfiguredSentry(DummySentry):
         def __init__(self):
             """Initialize the Sentry SDK."""
             ignore_logger("airflow.task")
-            executor_name = conf.get("core", "EXECUTOR")
 
             sentry_flask = FlaskIntegration()
 
             # LoggingIntegration is set by default.
             integrations = [sentry_flask]
 
-            if executor_name == "CeleryExecutor":
+            executor_class, _ = ExecutorLoader.import_default_executor_cls()
+
+            if executor_class.supports_sentry:
                 from sentry_sdk.integrations.celery import CeleryIntegration
 
                 sentry_celery = CeleryIntegration()

tests/executors/test_base_executor.py

@@ -32,6 +32,14 @@
 from airflow.utils.state import State
 
 
+def test_supports_sentry():
+    assert not BaseExecutor.supports_sentry
+
+
+def test_supports_pickling():
+    assert BaseExecutor.supports_pickling
+
+
 def test_is_local_default_value():
     assert not BaseExecutor.is_local
 

tests/executors/test_celery_executor.py

@@ -35,6 +35,7 @@
 
 from airflow.configuration import conf
 from airflow.executors import celery_executor
+from airflow.executors.celery_executor import CeleryExecutor
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.dag import DAG
 from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
@@ -98,6 +99,12 @@ def teardown_method(self) -> None:
         db.clear_db_runs()
         db.clear_db_jobs()
 
+    def test_supports_pickling(self):
+        assert CeleryExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert CeleryExecutor.supports_sentry
+
     @pytest.mark.quarantined
     @pytest.mark.backend("mysql", "postgres")
     def test_exception_propagation(self):

tests/executors/test_celery_kubernetes_executor.py

@@ -31,6 +31,12 @@
 
 
 class TestCeleryKubernetesExecutor:
+    def test_supports_pickling(self):
+        assert CeleryKubernetesExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert not CeleryKubernetesExecutor.supports_sentry
+
     def test_is_local_default_value(self):
         assert not CeleryKubernetesExecutor.is_local
 

tests/executors/test_dask_executor.py

@@ -90,6 +90,12 @@ def setup_method(self):
         self.dagbag = DagBag(include_examples=True)
         self.cluster = LocalCluster()
 
+    def test_supports_pickling(self):
+        assert not DaskExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert not DaskExecutor.supports_sentry
+
     def test_dask_executor_functions(self):
         executor = DaskExecutor(cluster_address=self.cluster.scheduler_address)
         self.assert_tasks_on_executor(executor, timeout_executor=120)

tests/executors/test_kubernetes_executor.py

@@ -1109,6 +1109,12 @@ def test_clear_not_launched_queued_tasks_clear_only_by_job_id(self, dag_maker, c
         assert ti0.state == State.SCHEDULED
         assert ti1.state == State.QUEUED
 
+    def test_supports_pickling(self):
+        assert KubernetesExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert not KubernetesExecutor.supports_sentry
+
 
 class TestKubernetesJobWatcher:
     test_namespace = "airflow"

tests/executors/test_local_executor.py

@@ -31,6 +31,12 @@ class TestLocalExecutor:
 
     TEST_SUCCESS_COMMANDS = 5
 
+    def test_supports_pickling(self):
+        assert not LocalExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert not LocalExecutor.supports_sentry
+
     def test_is_local_default_value(self):
         assert LocalExecutor.is_local
 

tests/executors/test_local_kubernetes_executor.py

@@ -26,6 +26,12 @@
 
 
 class TestLocalKubernetesExecutor:
+    def test_supports_pickling(self):
+        assert not LocalKubernetesExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert not LocalKubernetesExecutor.supports_sentry
+
     def test_is_local_default_value(self):
         assert not LocalKubernetesExecutor.is_local
 

tests/executors/test_sequential_executor.py

@@ -23,6 +23,12 @@
 
 
 class TestSequentialExecutor:
+    def test_supports_pickling(self):
+        assert not SequentialExecutor.supports_pickling
+
+    def test_supports_sentry(self):
+        assert not SequentialExecutor.supports_sentry
+
     def test_is_local_default_value(self):
         assert SequentialExecutor.is_local