how to specify celery queue dynamic by parmater before dag run

[SaithZhang]

Hi,
i use Airflow 2.0.1 with Celery worker in my production env.
The following is trigger by my java app with command like

/data/anaconda3/envs/airflow/bin/airflow dags trigger 'script_dag' -r `uuidgen` --conf '{"key1":"xxx", "key2":"script_id_123_job", "key3":"prod", "key4":"10.0.0.1"}'

My purpose is that set the key4 which is the machine ip to the queue param, because i want to the DAG execute
in the same ip with JAVA app .
I try

wonderful_curve = BashExtendOperator(
    task_id='wonderful_curve_task',
    bash_command='cd /data/ipva/ScriptEvaluation/ip_value/nlp &&  PYTHONIOENCODING=utf-8 '
                 '/data/anaconda3/envs/script_evaluation/bin/python -m '
                 'ScriptEvaluation.wonderful_curve \'{{ dag_run.conf["key1"] if dag_run else "" }}\' {{ dag_run.conf['
                 '"key2"] if dag_run else "" }} {{ dag_run.conf["key3"] if dag_run else "" }}',
    xcom_push=True,
   queue="{{ dag_run.conf["key4"] if dag_run else "" }}"
    dag=dag,
)

It doesn't work

i also add queue to my BashExtendOperator template_fields , doesn't work either.

template_fields = ('bash_command', 'env', 'queue')

And use socket.gethostbyname(socket.gethostname()) just give me the same ip which runs the scheduler instance,
but not the ip which i trigger the dag

wonderful_curve = BashExtendOperator(
    task_id='wonderful_curve_task',
    bash_command='cd /data/ipva/ScriptEvaluation/ip_value/nlp &&  PYTHONIOENCODING=utf-8 '
                 '/data/anaconda3/envs/script_evaluation/bin/python -m '
                 'ScriptEvaluation.wonderful_curve \'{{ dag_run.conf["key1"] if dag_run else "" }}\' {{ dag_run.conf['
                 '"key2"] if dag_run else "" }} {{ dag_run.conf["key3"] if dag_run else "" }}',
    xcom_push=True,
   queue='prod_script_' + socket.gethostbyname(socket.gethostname())"
    dag=dag,
)

FULL DAG

dag = DAG(
    dag_id='script_dag',
    default_args={"start_date": days_ago(2), "owner": "Airflow"},
    schedule_interval=None,
    dagrun_timeout=timedelta(minutes=60),
    tags=['algo']
)
...
wonderful_curve = BashExtendOperator(
    task_id='wonderful_curve_task',
    bash_command='cd /data/ipva/ScriptEvaluation/ip_value/nlp &&  PYTHONIOENCODING=utf-8 '
                 '/data/anaconda3/envs/script_evaluation/bin/python -m '
                 'ScriptEvaluation.wonderful_curve \'{{ dag_run.conf["key1"] if dag_run else "" }}\' {{ dag_run.conf['
                 '"key2"] if dag_run else "" }} {{ dag_run.conf["key3"] if dag_run else "" }}',
    xcom_push=True,
    queue='... set the queue by dynamic paramater ..."
    dag=dag,
)
...

run_analysis >> [..., wonderful_curve, ...]

Many thanks to your help!

[turbaszek]

Templating does not work as it happen at the moment of execution, so it's already on worker. I would suggest experimenting with cluster policies https://airflow.apache.org/docs/apache-airflow/stable/concepts.html#cluster-policy

def task_policy(task: BaseOperator):
  task.queue = task.dag.conf.get("key4", defaul_queue)

or if you use it for all task in DAG

def dag_policy(dag: DAG):
  queue = dag.conf.get("key4", default_queue)
  for task in dag.tasks:
     task.queue = queue

[SaithZhang]

Thanks for the quick reply.
I tried your way , but it seems airflow_local_settings.py not taking effect

airflow_local_settings.py

"""Airflow logging settings"""

from airflow.models import BaseOperator


def task_policy(task: BaseOperator):
    task.queue = task.dag.conf.get("key4", "defalut_test")

which i put it in /Users/saith/airflow/config

I see the scheduler log says that the airflow has loaded the airflow_local_settings.py

[2021-03-25 13:54:29,152] {scheduler_job.py:1834} INFO - Resetting orphaned tasks for active dag runs
[2021-03-25 13:54:29,158] {settings.py:54} INFO - Configured default timezone Timezone('UTC')
[2021-03-25 13:54:29,160] {settings.py:423} INFO - Loaded airflow_local_settings from /Users/saith/airflow/config/airflow_local_settings.py .

And i trigger the test_dag with

 airflow dags trigger 'dag_test' -r `uuidgen` --conf '{"key4":"test_param"}'

my test dag

from datetime import timedelta

from airflow.models import DAG
from airflow.utils.dates import days_ago
from my_operators.BashExtendOperator import BashExtendOperator

dag = DAG(
    dag_id='dag_test',
    default_args={"start_date": days_ago(2), "owner": "Airflow"},
    schedule_interval=None,
    dagrun_timeout=timedelta(minutes=60),
    tags=['algo']
)

run_analysis = BashExtendOperator(
    task_id='start_run_analysis_task',
    bash_command='echo \'123\' ',
    xcom_push=True,
    dag=dag,
)

extract_biography = BashExtendOperator(
    task_id='extract_biography_task',
    bash_command='echo \'234\' ',
    xcom_push=True,
    dag=dag,
)

run_analysis >> extract_biography

I check the Task Instance Details in webUI
and see that

queue        default

queue is still default.

I find the error log in scheduler log, but i still dont know how to get the conf through DAG object

[2021-03-25 14:24:48,735] {logging_mixin.py:104} INFO - [2021-03-25 14:24:48,735] {dagbag.py:475} ERROR - 'DAG' object has no attribute 'conf'
Traceback (most recent call last):
  File "/Users/saith/anaconda3/envs/airflow/lib/python3.7/site-packages/airflow/models/dagbag.py", line 462, in collect_dags
    found_dags = self.process_file(filepath, only_if_updated=only_if_updated, safe_mode=safe_mode)
  File "/Users/saith/anaconda3/envs/airflow/lib/python3.7/site-packages/airflow/models/dagbag.py", line 274, in process_file
    found_dags = self._process_modules(filepath, mods, file_last_changed_on_disk)
  File "/Users/saith/anaconda3/envs/airflow/lib/python3.7/site-packages/airflow/models/dagbag.py", line 374, in _process_modules
    self.bag_dag(dag=dag, root_dag=dag)
  File "/Users/saith/anaconda3/envs/airflow/lib/python3.7/site-packages/airflow/models/dagbag.py", line 401, in bag_dag
    settings.task_policy(task)
  File "/Users/saith/airflow/config/airflow_local_settings.py", line 24, in task_policy
    task.queue = task.dag.conf.get("key4", "defalut_test")
AttributeError: 'DAG' object has no attribute 'conf'
[2021-03-25 14:24:48,736] {scheduler_job.py:641} WARNING - No viable dags retrieved from /Users/saith/airflow/dags/dag_test.py
[2021-03-25 14:24:48,736] {scheduler_job.py:190} INFO - Processing /Users/saith/airflow/dags/dag_test.py took 0.015 seconds

Do you know how i can get the config through DAG object , thanks
I test it in my local mac.

[SaithZhang]

Hi, after search the issue #13811
I find the following will work in Airflow 2.0.1

def task_instance_mutation_hook(task_instance):
    dag_run = task_instance.get_dagrun()
    conf = dag_run.conf
    # print(conf)
    if conf['key4']:
        if 'change_queue' in conf:
            return
        else:
            task_instance.queue = conf['key4']
            conf["change_queue"] = 'true'

[turbaszek]

Right, the conf is attribute of DagRun not DAG. Good to know that you worked it out! 👏