Node label/capability based scheduling (for CeleryExecutor)

[hterik]

Description

Schedule tasks based on labels/capabilities with and+or matching, instead of just queue names.

Use case/motivation

Imagine you have tasks that can only run on specific nodes, for example GPUs, different operating systems or other external HW-peripherals connected.

Some tasks have partial constraints,
Task A: requires gpu
Task B: requires linux
Task C: requires gpu && linux
Task D: requires gpu && windows
Task E: requires gpu || mac-m1-cpu

To serve the above you might have 5 different nodes
Node 1: Linux (can serve task B only)
Node 2: Linux+Gpu (can serve A+B+C+D+E)
Node 3: Windows+Gpu (can serve A+D)

Optimizing this type of planning is, from what i understand, not possible with CeleryExecutor today. Airflow celery workers can only listen to a list of queues. In the scenario above, task D should not be assigned to a single queue, because there are two workers that could potentially execute it. Node2+3 can't listen to a common queue though because they support different features. Instead you would have to compromise and choose one of the nodes for task D.

Is this even possible to solve with Celery? I've scrutinized the underlying Celery documents deeply and it has lots of advanced features for inserting items to the queues using exchanges and routers, but it seems like consumers only read from fixed queue-names.

Am i missing something or is this correct? It feels like this should be a very common scenario.
Compared with other systems, Kubernetes can do it with nodeSelectors, affinity and anti-affinity. Using Airflow KubernetesExecutor these can be injected in pod_override.
Jenkins has label conditions

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[jscheffl]

Such feature is indirectly supported with a single "label" using the "queue" feature. You can assign tasks to "queues" and workers can pull work from one or multiple queues. (See https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/executor/celery.html#queues)

What is not currently possible is an && and || selection with multiple labels. But this is hard to achieve - at least in Celery - because the work is queues in Redis (if not configured otherwise) and the queue names are used for the Redis queue. Adding a logic of multiple labels would probably re-work a lot of logic.
But it is mainly inherited from the Celery-project. It is not specific to Airflow.

So I assume (from my background) you need to request the feature rather in Celery (--> https://github.com/celery/celery) and once Airflow upgrades the feature much probably propagates here.

Otherwise you can convince me of any other that this/such feature can be implemented in Airflow.

[jscheffl]

To map the current queue features to your demand, a workaround to implement the demand you are raising your be:

• Node 1: Listen to queue linux and x86
• Node 2: Listen to queue linux, macm1_gpu, gpu and linux_gpu
• Node 3: Listen to queue windows, gpu and windows_gpu - note that Celery does not support WIndows though :-D

Tasks can use labels:

• Task A: queue=gpu
• Task B: queue=linux
• Task C: queue=linux_gpu
• Task D: queue=windows_gpu
• Task E: queue=macm1_gpu

Can you let me know if under these exampldes the request can be closed - or do you open the request on Celery side?

[github-actions]

This issue has been automatically marked as stale because it has been open for 14 days with no response from the author. It will be closed in next 7 days if no further activity occurs from the issue author.

[github-actions]

This issue has been closed because it has not received response from the issue author.