Move callback_sink from executor to scheduler

airflow/executors/base_executor.py

@@ -19,8 +19,6 @@
 from collections import OrderedDict
 from typing import Any, Counter, Dict, List, Optional, Set, Tuple, Union
 
-from airflow.callbacks.base_callback_sink import BaseCallbackSink
-from airflow.callbacks.callback_requests import CallbackRequest
 from airflow.configuration import conf
 from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
 from airflow.stats import Stats
@@ -63,7 +61,6 @@ class BaseExecutor(LoggingMixin):
     """
 
     job_id: Union[None, int, str] = None
-    callback_sink: Optional[BaseCallbackSink] = None
 
     def __init__(self, parallelism: int = PARALLELISM):
         super().__init__()
@@ -350,14 +347,3 @@ def debug_dump(self):
             len(self.event_buffer),
             "\n\t".join(map(repr, self.event_buffer.items())),
         )
-
-    def send_callback(self, request: CallbackRequest) -> None:
-        """Sends callback for execution.
-
-        Provides a default implementation which sends the callback to the `callback_sink` object.
-
-        :param request: Callback request to be executed.
-        """
-        if not self.callback_sink:
-            raise ValueError("Callback sink is not ready.")
-        self.callback_sink.send(request)

airflow/executors/celery_kubernetes_executor.py

@@ -17,8 +17,6 @@
 # under the License.
 from typing import Dict, List, Optional, Set, Union
 
-from airflow.callbacks.base_callback_sink import BaseCallbackSink
-from airflow.callbacks.callback_requests import CallbackRequest
 from airflow.configuration import conf
 from airflow.executors.base_executor import CommandType, EventBufferValueType, QueuedTaskInstanceType
 from airflow.executors.celery_executor import CeleryExecutor
@@ -37,7 +35,6 @@ class CeleryKubernetesExecutor(LoggingMixin):
     """
 
     supports_ad_hoc_ti_run: bool = True
-    callback_sink: Optional[BaseCallbackSink] = None
 
     KUBERNETES_QUEUE = conf.get('celery_kubernetes_executor', 'kubernetes_queue')
 
@@ -207,12 +204,3 @@ def debug_dump(self) -> None:
         self.celery_executor.debug_dump()
         self.log.info("Dumping KubernetesExecutor state")
         self.kubernetes_executor.debug_dump()
-
-    def send_callback(self, request: CallbackRequest) -> None:
-        """Sends callback for execution.
-
-        :param request: Callback request to be executed.
-        """
-        if not self.callback_sink:
-            raise ValueError("Callback sink is not ready.")
-        self.callback_sink.send(request)

airflow/executors/local_kubernetes_executor.py

@@ -17,8 +17,6 @@
 # under the License.
 from typing import Dict, List, Optional, Set, Union
 
-from airflow.callbacks.base_callback_sink import BaseCallbackSink
-from airflow.callbacks.callback_requests import CallbackRequest
 from airflow.configuration import conf
 from airflow.executors.base_executor import CommandType, EventBufferValueType, QueuedTaskInstanceType
 from airflow.executors.kubernetes_executor import KubernetesExecutor
@@ -37,7 +35,6 @@ class LocalKubernetesExecutor(LoggingMixin):
     """
 
     supports_ad_hoc_ti_run: bool = True
-    callback_sink: Optional[BaseCallbackSink] = None
 
     KUBERNETES_QUEUE = conf.get('local_kubernetes_executor', 'kubernetes_queue')
 
@@ -206,12 +203,3 @@ def debug_dump(self) -> None:
         self.local_executor.debug_dump()
         self.log.info("Dumping KubernetesExecutor state")
         self.kubernetes_executor.debug_dump()
-
-    def send_callback(self, request: CallbackRequest) -> None:
-        """Sends callback for execution.
-
-        :param request: Callback request to be executed.
-        """
-        if not self.callback_sink:
-            raise ValueError("Callback sink is not ready.")
-        self.callback_sink.send(request)

airflow/jobs/scheduler_job.py

@@ -34,7 +34,13 @@
 from sqlalchemy.orm.session import Session, make_transient
 
 from airflow import models, settings
-from airflow.callbacks.callback_requests import DagCallbackRequest, SlaCallbackRequest, TaskCallbackRequest
+from airflow.callbacks.base_callback_sink import BaseCallbackSink
+from airflow.callbacks.callback_requests import (
+    CallbackRequest,
+    DagCallbackRequest,
+    SlaCallbackRequest,
+    TaskCallbackRequest,
+)
 from airflow.callbacks.database_callback_sink import DatabaseCallbackSink
 from airflow.callbacks.pipe_callback_sink import PipeCallbackSink
 from airflow.configuration import conf
@@ -165,6 +171,8 @@ def __init__(
                 DeprecationWarning,
             )
 
+        self._callback_sink: Optional[BaseCallbackSink] = None
+
     def register_signals(self) -> None:
         """Register signals that stop child processes"""
         signal.signal(signal.SIGINT, self._exit_gracefully)
@@ -687,7 +695,7 @@ def _process_executor_events(self, session: Session = None) -> int:
                         simple_task_instance=SimpleTaskInstance.from_ti(ti),
                         msg=msg % (ti, state, ti.state, info),
                     )
-                    self.executor.send_callback(request)
+                    self._send_callback(request)
                 else:
                     ti.handle_failure(error=msg % (ti, state, ti.state, info), session=session)
 
@@ -721,12 +729,10 @@ def _execute(self) -> None:
             self.executor.job_id = self.id
             if self.processor_agent:
                 self.log.debug("Using PipeCallbackSink as callback sink.")
-                self.executor.callback_sink = PipeCallbackSink(
-                    get_sink_pipe=self.processor_agent.get_callbacks_pipe
-                )
+                self._callback_sink = PipeCallbackSink(get_sink_pipe=self.processor_agent.get_callbacks_pipe)
             else:
                 self.log.debug("Using DatabaseCallbackSink as callback sink.")
-                self.executor.callback_sink = DatabaseCallbackSink()
+                self._callback_sink = DatabaseCallbackSink()
 
             self.executor.start()
 
@@ -1203,7 +1209,7 @@ def _verify_integrity_if_dag_changed(self, dag_run: DagRun, session=None):
     def _send_dag_callbacks_to_processor(self, dag: DAG, callback: Optional[DagCallbackRequest] = None):
         self._send_sla_callbacks_to_processor(dag)
         if callback:
-            self.executor.send_callback(callback)
+            self._send_callback(callback)
         else:
             self.log.debug("callback is empty")
 
@@ -1217,7 +1223,7 @@ def _send_sla_callbacks_to_processor(self, dag: DAG):
             return
 
         request = SlaCallbackRequest(full_filepath=dag.fileloc, dag_id=dag.dag_id)
-        self.executor.send_callback(request)
+        self._send_callback(request)
 
     @provide_session
     def _emit_pool_metrics(self, session: Session = None) -> None:
@@ -1378,5 +1384,14 @@ def _find_zombies(self, session):
                 msg=f"Detected {ti} as zombie",
             )
             self.log.error("Detected zombie job: %s", request)
-            self.executor.send_callback(request)
+            self._send_callback(request)
             Stats.incr('zombies_killed')
+
+    def _send_callback(self, request: CallbackRequest) -> None:
+        """Sends callback for execution.
+
+        :param request: Callback request to be executed.
+        """
+        if not self._callback_sink:
+            raise ValueError("Callback sink is not ready.")
+        self._callback_sink.send(request)

tests/executors/test_celery_kubernetes_executor.py

@@ -19,7 +19,6 @@
 
 from parameterized import parameterized
 
-from airflow.callbacks.callback_requests import CallbackRequest
 from airflow.configuration import conf
 from airflow.executors.celery_executor import CeleryExecutor
 from airflow.executors.celery_kubernetes_executor import CeleryKubernetesExecutor
@@ -224,14 +223,3 @@ def test_kubernetes_executor_knows_its_queue(self):
         assert k8s_executor_mock.kubernetes_queue == conf.get(
             'celery_kubernetes_executor', 'kubernetes_queue'
         )
-
-    def test_send_callback(self):
-        cel_exec = CeleryExecutor()
-        k8s_exec = KubernetesExecutor()
-        cel_k8s_exec = CeleryKubernetesExecutor(cel_exec, k8s_exec)
-        cel_k8s_exec.callback_sink = mock.MagicMock()
-
-        callback = CallbackRequest(full_filepath="fake")
-        cel_k8s_exec.send_callback(callback)
-
-        cel_k8s_exec.callback_sink.send.assert_called_once_with(callback)

tests/executors/test_local_kubernetes_executor.py

@@ -17,7 +17,6 @@
 # under the License.
 from unittest import mock
 
-from airflow.callbacks.callback_requests import CallbackRequest
 from airflow.configuration import conf
 from airflow.executors.local_executor import LocalExecutor
 from airflow.executors.local_kubernetes_executor import LocalKubernetesExecutor
@@ -68,14 +67,3 @@ def test_kubernetes_executor_knows_its_queue(self):
         LocalKubernetesExecutor(local_executor_mock, k8s_executor_mock)
 
         assert k8s_executor_mock.kubernetes_queue == conf.get('local_kubernetes_executor', 'kubernetes_queue')
-
-    def test_send_callback(self):
-        local_executor_mock = mock.MagicMock()
-        k8s_executor_mock = mock.MagicMock()
-        local_k8s_exec = LocalKubernetesExecutor(local_executor_mock, k8s_executor_mock)
-        local_k8s_exec.callback_sink = mock.MagicMock()
-
-        callback = CallbackRequest(full_filepath="fake")
-        local_k8s_exec.send_callback(callback)
-
-        local_k8s_exec.callback_sink.send.assert_called_once_with(callback)