Dynamic queue selection for heterogeneous workers

[MaicoTimmerman]

Our Airflow deployment uses heterogeneous workers, where we route tasks to respective workers. We have use-cases where we would dynamically determine the routing. Examples include:

• Directing traffic to workers located in different availability zones.
• Sending tasks to workers collocated with resources it depends on.

I've experimented with using triggers setting the TaskInstance.queue:

class SetQueueTriggerEvent(TriggerEvent):
    def __init__(self, queue: str) -> None:
        super().__init__(payload=queue)
        self.queue = queue

    @provide_session
    def handle_submit(self, *, task_instance: TaskInstance, session: Session = NEW_SESSION) -> None:
        super().handle_submit(task_instance=task_instance, session=session)
        logger.info(f"Setting task_instance {task_instance} queue to {self.queue}")
        task_instance.queue = self.queue  # type: ignore[assignment]  # Sqlalchemy 2.0 required (soon tm)

Unfortunately, when testing in practice, we run into issues as there multiple instance where TaskInstance.queue is refreshed from the actual DAG's task definition (airflow.models.taskinstance._refresh_from_task), taking the original value from serialized_dag table. Ultimately, to stop the worker/scheduler from putting back the original queue, I'd need to make sure that serialized_dag gets the new value (permanently).

Question 1: Is there an alternative way to dynamically select queues for task, or do I have to solve it in the DAG processor?

As potential solution I've been playing with the thought of creating a @property that shadows BaseOperator.queue, where we influence DAG parsing (and thus serialized_dag). However, with Airflow's DAG versioning from Airflow 3.x in mind I'm not to happy about creating count(DAGRun) entries in the serialized DAG table.

Question 2: Does modifying queue argument create new DAG versions in Airflow 3.x? I could not locate any information on DAG versioning in the documentation. Astronomer documentation on Dag versioning together with some initial code walkthorugh seems to suggest that the queue is included when evaluating new DAG versions, as it is part of the serialized_dag table.