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.

[zach-overflow]

I can provide some context for question 1. As for question 2, that's a very good question and I would also like to know the answer.

Question 1
Yes, you definitely will want to take a look at cluster policies. If I'm understanding your described use-case, then specifically either a task_policy or task_instance_mutation_hook, or possibly a combination of both would very likely suit your needs well. I have experience doing something similar with just a custom task_policy for dynamic queue assignment.

Question 2
I would also like to know the answer to this. My purely speculative guess is that the answer depends on what the given DAG Bundle versioning logic is, since my understanding of this doc is that DAG versioning behavior may differ depending on the underlying bundle class' implementation.

[potiuk]

Not sure about 2 either. Worth checking. But I think the right solution for this would be to built-in a dynamic queue capability in Airflow to do it differently than via cloud policy.

This however would require similar approach as custom timetables - likely this should require a dynamic "custom queue selector" object that you should pass instead of queue.

The issue here is that queue is deterrmined in scheduler before task is run (classic chicken-egg - we cannot run code in a worker to determine the queue, because we need to have a queue to run the task!). So the issue here is that - due to security model of Airflow - Dag Authors should not bve write any dynamic code to determine the queue. Such code has to be installed as a plugin by Deployment Manager (similarly as cluster policies are) - not DagAuthors.