[D] Dynamic queue selection for heterogeneous workers [airflow]

["MaicoTimmerman (via GitHub)" <gi...@apache.org>]

GitHub user MaicoTimmerman created a discussion: Dynamic queue selection for heterogeneous workers

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`:
```python
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](https://www.astronomer.io/docs/learn/airflow-dag-versioning#dag-versioning-vs-dag-bundles) 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.


GitHub link: https://github.com/apache/airflow/discussions/57477

----
This is an automatically sent email for commits@airflow.apache.org.
To unsubscribe, please send an email to: commits-unsubscribe@airflow.apache.org

Re: [D] Dynamic queue selection for heterogeneous workers [airflow]

["potiuk (via GitHub)" <gi...@apache.org>]

GitHub user potiuk added a comment to the discussion: Dynamic queue selection for heterogeneous workers

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.


GitHub link: https://github.com/apache/airflow/discussions/57477#discussioncomment-15304917

----
This is an automatically sent email for commits@airflow.apache.org.
To unsubscribe, please send an email to: commits-unsubscribe@airflow.apache.org

Re: [D] Dynamic queue selection for heterogeneous workers [airflow]

["zach-overflow (via GitHub)" <gi...@apache.org>]

GitHub user zach-overflow added a comment to the discussion: Dynamic queue selection for heterogeneous workers

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](https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/cluster-policies.html#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](http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html#other-considerations) is that DAG versioning behavior may differ depending on the underlying bundle class' implementation.

GitHub link: https://github.com/apache/airflow/discussions/57477#discussioncomment-15301096

----
This is an automatically sent email for commits@airflow.apache.org.
To unsubscribe, please send an email to: commits-unsubscribe@airflow.apache.org