Enable trigger queue support for async callbacks

Author: zach-overflow

airflow-core/docs/authoring-and-scheduling/deferring.rst

@@ -493,6 +493,8 @@ According to `benchmarks <https://github.com/apache/airflow/pull/58803#pullreque
 
 You can determine a suitable value for your deployment by creating a large number of triggers (for example, by triggering a Dag with many deferrable tasks) and observing both how the load is distributed across Triggerers in your environment and how long it takes for all Triggerers to pick up the triggers.
 
+.. _deferring/triggerer_queue_assignment:
+
 Controlling Triggerer Host Assignment Per Trigger
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -508,7 +510,7 @@ Under some circumstances, it may be desirable to assign a Trigger to a specific
     If you are using :doc:`Multi-Team mode</core-concepts/multi-team>`, the ``--team-name`` option provides
     native team-scoped triggerer assignment for all trigger types (task-created, event-driven, and callback).
     See :ref:`Team-scoped Triggerer <multi-team-triggerer>` in the Multi-Team documentation.
-    The ``--queues`` option described below is an older, queue-based mechanism that can be combined with
+    The ``--queues`` option described below is a distinct queue-based mechanism which may be combined with
     ``--team-name`` if needed.
 
 To enable queue assignment for triggers, do the following:
@@ -536,33 +538,35 @@ This feature is only compatible with executors which utilize the task ``queue``
 task-created triggers. :doc:`Event-Driven Triggers<../authoring-and-scheduling/event-scheduling>`
 are not tied to a task queue, but a :class:`~airflow.triggers.base.BaseEventTrigger` subclass can
 be assigned to a queue explicitly by passing ``queue=`` to ``super().__init__()``, independently
-of any executor.
-
-Queue assignment is not supported for triggers from async
-:doc:`Callbacks<../administration-and-deployment/logging-monitoring/callbacks>`.
-
-+------------------------------------------------------------------------------------------------------+----------------------+----------------------------------------------------------------------------------+
-| Trigger Type                                                                                         | Supports queues?     | Triggerer assignment when :ref:`config:triggerer__queues_enabled` is ``True``    |
-+======================================================================================================+======================+==================================================================================+
-| Task-created Trigger instances                                                                       | Yes                  | Any triggerer with the task queue present in its ``--queues`` option             |
-+------------------------------------------------------------------------------------------------------+----------------------+----------------------------------------------------------------------------------+
-| :doc:`Event-Driven Triggers<../authoring-and-scheduling/event-scheduling>`                           | Yes                  | Any triggerer whose ``--queues`` includes the trigger's ``queue``, or any        |
-|                                                                                                      |                      | triggerer without ``--queues`` if no ``queue`` was set                           |
-+------------------------------------------------------------------------------------------------------+----------------------+----------------------------------------------------------------------------------+
-| Triggers from async :doc:`Callbacks<../administration-and-deployment/logging-monitoring/callbacks>`  | No                   | Any triggerer running without the ``--queues`` option                            |
-+------------------------------------------------------------------------------------------------------+----------------------+----------------------------------------------------------------------------------+
-
-If you use queues for task-based and/or event-driven triggers, while **also** using callback triggers
-(or event-driven triggers without an explicit ``queue``), you must run one or more triggerer hosts
+of any executor. Similarly, triggers from async
+:doc:`Callbacks<../administration-and-deployment/logging-monitoring/callbacks>` can be assigned to a
+queue by passing ``queue=`` to :class:`~airflow.sdk.AsyncCallback`.
+
++------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------+
+| Trigger Type                                                                                         | Triggerer assignment source when :ref:`config:triggerer__queues_enabled` is ``True``    |
++======================================================================================================+=========================================================================================+
+| Task-created Trigger instances                                                                       | Any triggerer with the task queue present in its ``--queues`` option                    |
++------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------+
+| :doc:`Event-Driven Triggers<../authoring-and-scheduling/event-scheduling>`                           | Any triggerer whose ``--queues`` includes the trigger's ``queue``, or any               |
+|                                                                                                      | triggerer without ``--queues`` if no ``queue`` was set                                  |
++------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------+
+| Triggers from async :doc:`Callbacks<../administration-and-deployment/logging-monitoring/callbacks>`  | Any triggerer whose ``--queues`` includes the callback's ``queue``, or any              |
+|                                                                                                      | triggerer without ``--queues`` if no ``queue`` was set                                  |
++------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------+
+
+If you run a mix of triggers with queues **and** without queues assigned, (of any trigger type), you must run at least one triggerer host
 **without** the ``--queues`` option, so the latter are still run.
 
 .. note::
     To enable trigger queues, you must set the ``--queues`` option on one or more triggerers' startup command (these values may differ between the various triggerers).
-    If you set the ``--queue`` value of a triggerer to some value which no task queues or event-driven trigger queues exist for, that triggerer will never run any triggers.
-    Similarly, all ``triggerer`` instances running without the ``--queues`` option will only consume callback-based triggers and event-driven triggers that were not assigned an explicit ``queue``.
+    If you set the ``--queue`` value of a triggerer to some value which no triggers are assigned (either from their task queue, or from their explicit ``queue`` field), then that triggerer
+    will never run any triggers. Similarly, all triggerer instances running without the ``--queues`` option will only consume callback-based triggers and event-driven triggers that were not
+    assigned an explicit ``queue``.
 
 An event-driven trigger's ``queue`` is set once, when the trigger is registered for an
-:class:`~airflow.sdk.AssetWatcher` during Dag processing, and is independent of
+:class:`~airflow.sdk.AssetWatcher` during Dag processing. A callback trigger's ``queue`` is set
+from the ``queue`` passed to :class:`~airflow.sdk.AsyncCallback` when the deadline is missed and the
+callback is queued on the Triggerer. Both are independent of
 :doc:`Multi-Team mode</core-concepts/multi-team>`'s ``team_name`` scoping — the two can be
 combined.
 

airflow-core/docs/howto/deadline-alerts.rst

@@ -362,6 +362,21 @@ A **custom asynchronous callback** might look like this:
     ):
         EmptyOperator(task_id="example_task")
 
+.. tip::
+    ``AsyncCallback`` accepts an optional ``queue`` parameter to assign the resulting trigger to a
+    specific :ref:`triggerer queue <config:triggerer__queues_enabled>`. If not specified, the
+    callback trigger runs on any triggerer started without the ``--queues`` option. See
+    :ref:`Controlling Triggerer Host Assignment Per Trigger <deferring/triggerer_queue_assignment>`
+    for details.
+
+    .. code-block:: python
+
+        AsyncCallback(
+            my_callback,
+            kwargs={"msg": "deadline missed"},
+            queue="alerts",
+        )
+
 Templating and Context
 ^^^^^^^^^^^^^^^^^^^^^^
 

airflow-core/src/airflow/models/callback.py

@@ -102,6 +102,13 @@ class ImportPathExecutorCallbackDefProtocol(ImportPathCallbackDefProtocol, Proto
     executor: str | None
 
 
+@runtime_checkable
+class ImportPathAsyncCallbackDefProtocol(ImportPathCallbackDefProtocol, Protocol):
+    """Protocol for callbacks that use the import path fetch method and support triggerer queue assignment."""
+
+    queue: str | None
+
+
 class Callback(Base, BaseWorkload):
     """Base class for callbacks."""
 
@@ -201,7 +208,7 @@ def create_from_sdk_def(callback_def: CallbackDefinitionProtocol, **kwargs) -> C
         match type(callback_def).__name__:
             case "AsyncCallback":
                 if TYPE_CHECKING:
-                    assert isinstance(callback_def, ImportPathCallbackDefProtocol)
+                    assert isinstance(callback_def, ImportPathAsyncCallbackDefProtocol)
                 return TriggererCallback(callback_def, **kwargs)
 
             case "SyncCallback":
@@ -244,6 +251,7 @@ def queue(self, *, session: Session) -> None:
             CallbackTrigger(
                 callback_path=self.data["path"],
                 callback_kwargs=self.data["kwargs"],
+                queue=self.data.get("queue"),
             )
         )
         self.trigger.team_name = team_name

airflow-core/src/airflow/models/trigger.py

@@ -200,7 +200,7 @@ def rotate_fernet_key(self):
     def from_object(cls, trigger: BaseTrigger) -> Trigger:
         """Alternative constructor that creates a trigger row based directly off of a Trigger object."""
         classpath, kwargs = trigger.serialize()
-        return cls(classpath=classpath, kwargs=kwargs)
+        return cls(classpath=classpath, kwargs=kwargs, queue=getattr(trigger, "queue", None))
 
     @classmethod
     @provide_session

airflow-core/src/airflow/triggers/callback.py

@@ -35,12 +35,20 @@
 class CallbackTrigger(BaseTrigger):
     """Trigger that executes a callback function asynchronously."""
 
-    supports_triggerer_queue: bool = False
+    supports_triggerer_queue: bool = True
 
-    def __init__(self, callback_path: str, callback_kwargs: dict[str, Any] | None = None):
+    def __init__(
+        self,
+        callback_path: str,
+        callback_kwargs: dict[str, Any] | None = None,
+        *,
+        queue: str | None = None,
+    ):
         super().__init__()
         self.callback_path = callback_path
         self.callback_kwargs = callback_kwargs or {}
+        # Read by Trigger.from_object() when persisting the row; unused once running.
+        self.queue = queue
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (

airflow-core/tests/unit/models/test_callback.py

@@ -117,6 +117,15 @@ def test_get_metric_info(self):
             "dag_id": TEST_DAG_ID,
         }
 
+    def test_get_metric_info_includes_queue(self):
+        """``queue`` is stored in ``self.data`` for a queued AsyncCallback, so it flows into metric tags."""
+        queued_callback = AsyncCallback(async_callback, kwargs=TEST_CALLBACK_KWARGS, queue="custom-queue")
+        callback = TriggererCallback(queued_callback, prefix="deadline_alerts", dag_id=TEST_DAG_ID)
+
+        metric_info = callback.get_metric_info(CallbackState.SUCCESS, "0")
+
+        assert metric_info["tags"]["queue"] == "custom-queue"
+
     def test_get_metric_info_dict_values_are_stringified(self):
         """
         Regression for ``TypeError: unhashable type: 'dict'`` raised by OpenTelemetry's
@@ -185,8 +194,17 @@ def test_queue(self, session):
         assert isinstance(callback.trigger, Trigger)
         assert callback.trigger.kwargs["callback_path"] == TEST_ASYNC_CALLBACK.path
         assert callback.trigger.kwargs["callback_kwargs"] == TEST_ASYNC_CALLBACK.kwargs
+        assert callback.trigger.queue is None
         assert callback.state == CallbackState.QUEUED
 
+    def test_queue_populates_trigger_queue(self, session):
+        queued_callback = AsyncCallback(async_callback, kwargs=TEST_CALLBACK_KWARGS, queue="custom-queue")
+        callback = TriggererCallback(queued_callback)
+
+        callback.queue(session=session)
+
+        assert callback.trigger.queue == "custom-queue"
+
     @staticmethod
     def _queue_callback(session, *, has_bundle, has_team):
         from airflow.models.dagbundle import DagBundleModel

airflow-core/tests/unit/models/test_trigger.py

@@ -1107,6 +1107,24 @@ def test_serialize_sensitive_kwargs():
     assert "value2" not in trigger_row.encrypted_kwargs
 
 
+def test_from_object_reads_queue_from_trigger():
+    """A ``queue`` attribute on the trigger object is carried onto the persisted row."""
+    trigger_instance = SensitiveKwargsTrigger(param1="value1", param2="value2")
+    trigger_instance.queue = "custom-queue"
+
+    trigger_row: Trigger = Trigger.from_object(trigger_instance)
+
+    assert trigger_row.queue == "custom-queue"
+
+
+def test_from_object_defaults_queue_to_none_when_trigger_has_no_queue_attribute():
+    trigger_instance = SensitiveKwargsTrigger(param1="value1", param2="value2")
+
+    trigger_row: Trigger = Trigger.from_object(trigger_instance)
+
+    assert trigger_row.queue is None
+
+
 def test_kwargs_not_encrypted():
     """
     Tests that we don't decrypt kwargs if they aren't encrypted.

airflow-core/tests/unit/triggers/test_callback.py

@@ -78,6 +78,24 @@ def test_serialization(self, callback_init_kwargs, expected_serialized_kwargs):
             "callback_kwargs": expected_serialized_kwargs,
         }
 
+    def test_supports_triggerer_queue(self):
+        assert CallbackTrigger.supports_triggerer_queue is True
+
+    def test_queue_attribute_is_not_part_of_serialized_kwargs(self):
+        """``queue`` is read directly off the trigger by ``Trigger.from_object``, not via serialize()."""
+        trigger = CallbackTrigger(
+            callback_path=TEST_CALLBACK_PATH,
+            callback_kwargs=None,
+            queue="custom-queue",
+        )
+
+        assert trigger.queue == "custom-queue"
+        _, kwargs = trigger.serialize()
+        assert "queue" not in kwargs
+
+    def test_queue_defaults_to_none(self, trigger):
+        assert trigger.queue is None
+
     @pytest.mark.asyncio
     async def test_run_success_with_async_function(self, trigger, mock_import_string):
         """Test trigger handles async functions correctly."""

task-sdk/src/airflow/sdk/definitions/callback.py

@@ -151,16 +151,31 @@ class AsyncCallback(Callback):
     triggerer.
 
     It will be called with Airflow context and specified kwargs when a deadline is missed.
+
+    Pass ``queue`` to assign the resulting trigger to a specific triggerer queue (see
+    :ref:`config:triggerer__queues_enabled` and the ``--queues`` option of ``airflow triggerer``).
     """
 
-    def __init__(self, callback_callable: Callable | str, kwargs: dict | None = None):
+    queue: str | None
+
+    def __init__(
+        self,
+        callback_callable: Callable | str,
+        kwargs: dict | None = None,
+        queue: str | None = None,
+    ):
         super().__init__(callback_callable=callback_callable, kwargs=kwargs)
+        self.queue = queue
 
     @classmethod
     def verify_callable(cls, callback: Callable):
         if not (inspect.iscoroutinefunction(callback) or hasattr(callback, "__await__")):
             raise AttributeError(f"Provided callback {callback} is not awaitable.")
 
+    @classmethod
+    def serialized_fields(cls) -> tuple[str, ...]:
+        return super().serialized_fields() + ("queue",)
+
 
 class SyncCallback(Callback):
     """

task-sdk/tests/task_sdk/definitions/test_callback.py

@@ -240,12 +240,28 @@ def test_init_error(self):
         with pytest.raises(AttributeError, match="is not awaitable."):
             AsyncCallback(empty_sync_callback_for_deadline_tests)
 
+    @pytest.mark.parametrize(
+        "queue",
+        [pytest.param("custom-queue", id="with_queue"), pytest.param(None, id="without_queue")],
+    )
+    def test_init_queue(self, queue):
+        callback = AsyncCallback(TEST_CALLBACK_PATH, kwargs=TEST_CALLBACK_KWARGS, queue=queue)
+        assert callback.queue == queue
+
     def test_serialize_deserialize(self):
         callback = AsyncCallback(TEST_CALLBACK_PATH, kwargs=TEST_CALLBACK_KWARGS)
         serialized = serialize(callback)
         deserialized = cast("Callback", deserialize(serialized.copy()))
         assert callback == deserialized
 
+    def test_serialize_deserialize_round_trip_keeps_queue(self):
+        callback = AsyncCallback(TEST_CALLBACK_PATH, kwargs=TEST_CALLBACK_KWARGS, queue="custom-queue")
+
+        deserialized = cast("AsyncCallback", deserialize(serialize(callback)))
+
+        assert deserialized == callback
+        assert deserialized.queue == "custom-queue"
+
 
 class TestSyncCallback:
     @pytest.mark.parametrize(