Executor Synchronous callback workload

airflow-core/src/airflow/executors/base_executor.py

@@ -32,7 +32,9 @@
 from airflow.configuration import conf
 from airflow.executors import workloads
 from airflow.executors.executor_loader import ExecutorLoader
+from airflow.executors.workloads.task import TaskInstanceDTO
 from airflow.models import Log
+from airflow.models.callback import CallbackKey
 from airflow.observability.metrics import stats_utils
 from airflow.observability.trace import Trace
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -52,6 +54,7 @@
     from airflow.callbacks.callback_requests import CallbackRequest
     from airflow.cli.cli_config import GroupCommand
     from airflow.executors.executor_utils import ExecutorName
+    from airflow.executors.workloads.types import WorkloadKey
     from airflow.models.taskinstance import TaskInstance
     from airflow.models.taskinstancekey import TaskInstanceKey
 
@@ -143,6 +146,7 @@ class BaseExecutor(LoggingMixin):
     active_spans = ThreadSafeDict()
 
     supports_ad_hoc_ti_run: bool = False
+    supports_callbacks: bool = False
     supports_multi_team: bool = False
     sentry_integration: str = ""
 
@@ -186,8 +190,9 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
         self.parallelism: int = parallelism
         self.team_name: str | None = team_name
         self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {}
-        self.running: set[TaskInstanceKey] = set()
-        self.event_buffer: dict[TaskInstanceKey, EventBufferValueType] = {}
+        self.queued_callbacks: dict[str, workloads.ExecuteCallback] = {}
+        self.running: set[WorkloadKey] = set()
+        self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {}
         self._task_event_logs: deque[Log] = deque()
         self.conf = ExecutorConf(team_name)
 
@@ -203,7 +208,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
         :meta private:
         """
 
-        self.attempts: dict[TaskInstanceKey, RunningRetryAttemptType] = defaultdict(RunningRetryAttemptType)
+        self.attempts: dict[WorkloadKey, RunningRetryAttemptType] = defaultdict(RunningRetryAttemptType)
 
     def __repr__(self):
         _repr = f"{self.__class__.__name__}(parallelism={self.parallelism}"
@@ -224,10 +229,47 @@ def log_task_event(self, *, event: str, extra: str, ti_key: TaskInstanceKey):
         self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra))
 
     def queue_workload(self, workload: workloads.All, session: Session) -> None:
-        if not isinstance(workload, workloads.ExecuteTask):
-            raise ValueError(f"Un-handled workload kind {type(workload).__name__!r} in {type(self).__name__}")
-        ti = workload.ti
-        self.queued_tasks[ti.key] = workload
+        if isinstance(workload, workloads.ExecuteTask):
+            ti = workload.ti
+            self.queued_tasks[ti.key] = workload
+        elif isinstance(workload, workloads.ExecuteCallback):
+            if not self.supports_callbacks:
+                raise NotImplementedError(
+                    f"{type(self).__name__} does not support ExecuteCallback workloads. "
+                    f"Set supports_callbacks = True and implement callback handling in _process_workloads(). "
+                    f"See LocalExecutor or CeleryExecutor for reference implementation."
+                )
+            self.queued_callbacks[workload.callback.id] = workload
+        else:
+            raise ValueError(
+                f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. "
+                f"Workload must be one of: ExecuteTask, ExecuteCallback."
+            )
+
+    def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, workloads.All]]:
+        """
+        Select and return the next batch of workloads to schedule, respecting priority policy.
+
+        Priority Policy: Callbacks are scheduled before tasks (callbacks complete existing work).
+        Callbacks are processed in FIFO order. Tasks are sorted by priority_weight (higher priority first).
+
+        :param open_slots: Number of available execution slots
+        """
+        workloads_to_schedule: list[tuple[WorkloadKey, workloads.All]] = []
+
+        if self.queued_callbacks:
+            for key, workload in self.queued_callbacks.items():
+                if len(workloads_to_schedule) >= open_slots:
+                    break
+                workloads_to_schedule.append((key, workload))
+
+        if open_slots > len(workloads_to_schedule) and self.queued_tasks:
+            for task_key, task_workload in self.order_queued_tasks_by_priority():
+                if len(workloads_to_schedule) >= open_slots:
+                    break
+                workloads_to_schedule.append((task_key, task_workload))
+
+        return workloads_to_schedule
 
     def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
         """
@@ -266,10 +308,10 @@ def heartbeat(self) -> None:
         """Heartbeat sent to trigger new jobs."""
         open_slots = self.parallelism - len(self.running)
 
-        num_running_tasks = len(self.running)
-        num_queued_tasks = len(self.queued_tasks)
+        num_running_workloads = len(self.running)
+        num_queued_workloads = len(self.queued_tasks) + len(self.queued_callbacks)
 
-        self._emit_metrics(open_slots, num_running_tasks, num_queued_tasks)
+        self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
         self.trigger_tasks(open_slots)
 
         # Calling child class sync method
@@ -350,16 +392,16 @@ def order_queued_tasks_by_priority(self) -> list[tuple[TaskInstanceKey, workload
 
     def trigger_tasks(self, open_slots: int) -> None:
         """
-        Initiate async execution of the queued tasks, up to the number of available slots.
+        Initiate async execution of queued workloads (tasks and callbacks), up to the number of available slots.
+
+        Callbacks are prioritized over tasks to complete existing work before starting new work.
 
         :param open_slots: Number of open slots
         """
-        sorted_queue = self.order_queued_tasks_by_priority()
+        workloads_to_schedule = self._get_workloads_to_schedule(open_slots)
         workload_list = []
 
-        for _ in range(min((open_slots, len(self.queued_tasks)))):
-            key, item = sorted_queue.pop()
-
+        for key, workload in workloads_to_schedule:
             # If a task makes it here but is still understood by the executor
             # to be running, it generally means that the task has been killed
             # externally and not yet been marked as failed.
@@ -373,12 +415,12 @@ def trigger_tasks(self, open_slots: int) -> None:
             if key in self.attempts:
                 del self.attempts[key]
 
-            if isinstance(item, workloads.ExecuteTask) and hasattr(item, "ti"):
-                ti = item.ti
+            if isinstance(workload, workloads.ExecuteTask) and hasattr(workload, "ti"):
+                ti = workload.ti
 
                 # If it's None, then the span for the current id hasn't been started.
                 if self.active_spans is not None and self.active_spans.get("ti:" + str(ti.id)) is None:
-                    if isinstance(ti, workloads.TaskInstance):
+                    if isinstance(ti, TaskInstanceDTO):
                         parent_context = Trace.extract(ti.parent_context_carrier)
                     else:
                         parent_context = Trace.extract(ti.dag_run.context_carrier)
@@ -397,7 +439,8 @@ def trigger_tasks(self, open_slots: int) -> None:
                     carrier = Trace.inject()
                     ti.context_carrier = carrier
 
-                workload_list.append(item)
+            workload_list.append(workload)
+
         if workload_list:
             self._process_workloads(workload_list)
 
@@ -459,24 +502,25 @@ def running_state(self, key: TaskInstanceKey, info=None) -> None:
         """
         self.change_state(key, TaskInstanceState.RUNNING, info, remove_running=False)
 
-    def get_event_buffer(self, dag_ids=None) -> dict[TaskInstanceKey, EventBufferValueType]:
+    def get_event_buffer(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueType]:
         """
         Return and flush the event buffer.
 
         In case dag_ids is specified it will only return and flush events
         for the given dag_ids. Otherwise, it returns and flushes all events.
+        Note: Callback events (with string keys) are always included regardless of dag_ids filter.
 
         :param dag_ids: the dag_ids to return events for; returns all if given ``None``.
         :return: a dict of events
         """
-        cleared_events: dict[TaskInstanceKey, EventBufferValueType] = {}
+        cleared_events: dict[WorkloadKey, EventBufferValueType] = {}
         if dag_ids is None:
             cleared_events = self.event_buffer
             self.event_buffer = {}
         else:
-            for ti_key in list(self.event_buffer.keys()):
-                if ti_key.dag_id in dag_ids:
-                    cleared_events[ti_key] = self.event_buffer.pop(ti_key)
+            for key in list(self.event_buffer.keys()):
+                if isinstance(key, CallbackKey) or key.dag_id in dag_ids:
+                    cleared_events[key] = self.event_buffer.pop(key)
 
         return cleared_events
 
@@ -529,21 +573,26 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
 
     @property
     def slots_available(self):
-        """Number of new tasks this executor instance can accept."""
-        return self.parallelism - len(self.running) - len(self.queued_tasks)
+        """Number of new workloads (tasks and callbacks) this executor instance can accept."""
+        return self.parallelism - len(self.running) - len(self.queued_tasks) - len(self.queued_callbacks)
 
     @property
     def slots_occupied(self):
-        """Number of tasks this executor instance is currently managing."""
-        return len(self.running) + len(self.queued_tasks)
+        """Number of workloads (tasks and callbacks) this executor instance is currently managing."""
+        return len(self.running) + len(self.queued_tasks) + len(self.queued_callbacks)
 
     def debug_dump(self):
         """Get called in response to SIGUSR2 by the scheduler."""
         self.log.info(
-            "executor.queued (%d)\n\t%s",
+            "executor.queued_tasks (%d)\n\t%s",
             len(self.queued_tasks),
             "\n\t".join(map(repr, self.queued_tasks.items())),
         )
+        self.log.info(
+            "executor.queued_callbacks (%d)\n\t%s",
+            len(self.queued_callbacks),
+            "\n\t".join(map(repr, self.queued_callbacks.items())),
+        )
         self.log.info("executor.running (%d)\n\t%s", len(self.running), "\n\t".join(map(repr, self.running)))
         self.log.info(
             "executor.event_buffer (%d)\n\t%s",

airflow-core/src/airflow/executors/local_executor.py

@@ -37,7 +37,8 @@
 
 from airflow.executors import workloads
 from airflow.executors.base_executor import BaseExecutor
-from airflow.utils.state import TaskInstanceState
+from airflow.executors.workloads.callback import execute_callback_workload
+from airflow.utils.state import CallbackState, TaskInstanceState
 
 # add logger to parameter of setproctitle to support logging
 if sys.platform == "darwin":
@@ -50,13 +51,23 @@
 if TYPE_CHECKING:
     from structlog.typing import FilteringBoundLogger as Logger
 
-    TaskInstanceStateType = tuple[workloads.TaskInstance, TaskInstanceState, Exception | None]
+    from airflow.executors.workloads.types import WorkloadResultType
+
+
+def _get_executor_process_title_prefix(team_name: str | None) -> str:
+    """
+    Build the process title prefix for LocalExecutor workers.
+
+    :param team_name: Team name from executor configuration
+    """
+    team_suffix = f" [{team_name}]" if team_name else ""
+    return f"airflow worker -- LocalExecutor{team_suffix}:"
 
 
 def _run_worker(
     logger_name: str,
     input: SimpleQueue[workloads.All | None],
-    output: Queue[TaskInstanceStateType],
+    output: Queue[WorkloadResultType],
     unread_messages: multiprocessing.sharedctypes.Synchronized[int],
     team_conf,
 ):
@@ -68,11 +79,8 @@ def _run_worker(
     log = structlog.get_logger(logger_name)
     log.info("Worker starting up pid=%d", os.getpid())
 
-    # Create team suffix for process title
-    team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else ""
-
     while True:
-        setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: <idle>", log)
+        setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} <idle>", log)
         try:
             workload = input.get()
         except EOFError:
@@ -87,25 +95,30 @@ def _run_worker(
             # Received poison pill, no more tasks to run
             return
 
-        if not isinstance(workload, workloads.ExecuteTask):
-            raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}")
-
         # Decrement this as soon as we pick up a message off the queue
         with unread_messages:
             unread_messages.value -= 1
-        key = None
-        if ti := getattr(workload, "ti", None):
-            key = ti.key
-        else:
-            raise TypeError(f"Don't know how to get ti key from {type(workload).__name__}")
 
-        try:
-            _execute_work(log, workload, team_conf)
+        # Handle different workload types
+        if isinstance(workload, workloads.ExecuteTask):
+            try:
+                _execute_work(log, workload, team_conf)
+                output.put((workload.ti.key, TaskInstanceState.SUCCESS, None))
+            except Exception as e:
+                log.exception("Task execution failed.")
+                output.put((workload.ti.key, TaskInstanceState.FAILED, e))
+
+        elif isinstance(workload, workloads.ExecuteCallback):
+            output.put((workload.callback.id, CallbackState.RUNNING, None))
+            try:
+                _execute_callback(log, workload, team_conf)
+                output.put((workload.callback.id, CallbackState.SUCCESS, None))
+            except Exception as e:
+                log.exception("Callback execution failed")
+                output.put((workload.callback.id, CallbackState.FAILED, e))
 
-            output.put((key, TaskInstanceState.SUCCESS, None))
-        except Exception as e:
-            log.exception("uhoh")
-            output.put((key, TaskInstanceState.FAILED, e))
+        else:
+            raise ValueError(f"LocalExecutor does not know how to handle {type(workload)}")
 
 
 def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> None:
@@ -118,9 +131,7 @@ def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> No
     """
     from airflow.sdk.execution_time.supervisor import supervise
 
-    # Create team suffix for process title
-    team_suffix = f" [{team_conf.team_name}]" if team_conf.team_name else ""
-    setproctitle(f"airflow worker -- LocalExecutor{team_suffix}: {workload.ti.id}", log)
+    setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.ti.id}", log)
 
     base_url = team_conf.get("api", "base_url", fallback="/")
     # If it's a relative URL, use localhost:8080 as the default
@@ -141,6 +152,22 @@ def _execute_work(log: Logger, workload: workloads.ExecuteTask, team_conf) -> No
     )
 
 
+def _execute_callback(log: Logger, workload: workloads.ExecuteCallback, team_conf) -> None:
+    """
+    Execute a callback workload.
+
+    :param log: Logger instance
+    :param workload: The ExecuteCallback workload to execute
+    :param team_conf: Team-specific executor configuration
+    """
+    setproctitle(f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.callback.id}", log)
+
+    success, error_msg = execute_callback_workload(workload.callback, log)
+
+    if not success:
+        raise RuntimeError(error_msg or "Callback execution failed")
+
+
 class LocalExecutor(BaseExecutor):
     """
     LocalExecutor executes tasks locally in parallel.
@@ -155,9 +182,10 @@ class LocalExecutor(BaseExecutor):
 
     supports_multi_team: bool = True
     serve_logs: bool = True
+    supports_callbacks: bool = True
 
     activity_queue: SimpleQueue[workloads.All | None]
-    result_queue: SimpleQueue[TaskInstanceStateType]
+    result_queue: SimpleQueue[WorkloadResultType]
     workers: dict[int, multiprocessing.Process]
     _unread_messages: multiprocessing.sharedctypes.Synchronized[int]
 
@@ -300,10 +328,14 @@ def end(self) -> None:
     def terminate(self):
         """Terminate the executor is not doing anything."""
 
-    def _process_workloads(self, workloads):
-        for workload in workloads:
+    def _process_workloads(self, workload_list):
+        for workload in workload_list:
             self.activity_queue.put(workload)
-            del self.queued_tasks[workload.ti.key]
+            # Remove from appropriate queue based on workload type
+            if isinstance(workload, workloads.ExecuteTask):
+                del self.queued_tasks[workload.ti.key]
+            elif isinstance(workload, workloads.ExecuteCallback):
+                del self.queued_callbacks[workload.callback.id]
         with self._unread_messages:
-            self._unread_messages.value += len(workloads)
+            self._unread_messages.value += len(workload_list)
         self._check_workers()