Task Group中使用Airflow Branch Operator报无效Task ID错误

from airflow.models import DAG
from airflow.operators.dummy import DummyOperator
from airflow.operators.python import BranchPythonOperator
from airflow.providers.mysql.operators.mysql import MySqlOperator
from typing import Optional

def skip_update_job_pod_name(dag):
    """
    :param dag: Airflow DAG
    :return: Dummy operator to skip update pod name
    """
    return DummyOperator(task_id="skip_update_job_pod_name", dag=dag)


def update_pod_name_branch_operator(dag: DAG, job_id: str):
    """branch operator to update pod name."""
    return BranchPythonOperator(
        dag=dag,
        trigger_rule="all_done",
        task_id="update_pod_name",
        python_callable=update_pod_name_func,
        op_kwargs={"job_id": job_id},
    )


def update_pod_name_func(job_id: Optional[str]) -> str:
    """function for update pod name."""
    return ["update_job_pod_name"] if job_id else ["skip_update_pod_name"]


def update_job_pod_name(dag: DAG, job_id: str, process_name: str) -> MySqlOperator:
    """
    :param dag: Airflow DAG
    :param job_id: Airflow Job ID
    :param process_name: name of the current running process
    :return: MySqlOperator to update Airflow job ID
    """
    return MySqlOperator(
        task_id="update_job_pod_name",
        mysql_conn_id="semantic-search-airflow-sdk",
        autocommit=True,
        sql=[
            f"""
                INSERT INTO airflow.Pod (job_id, pod_name, task_name)
                SELECT * FROM (SELECT '{job_id}', '{xcom_pull("pod_name")}', '{process_name}') AS temp
                WHERE NOT EXISTS (
                    SELECT pod_name FROM airflow.Pod WHERE pod_name = '{{{{ ti.xcom_pull(key="pod_name") }}}}'
                ) LIMIT 1;
            """
        ],
        task_concurrency=1,
        dag=dag,
        trigger_rule="all_done",
    )

airflow.exceptions.AirflowException: Branch callable must return valid task_ids. Invalid tasks found: {'update_job_pod_name'}

这个报错的核心原因是TaskGroup的task_id自动前缀规则和分支函数硬编码task_id不匹配：

• 当把Operator放到TaskGroup内部时，Airflow会自动给组内所有任务的task_id拼接前缀，最终生效的task_id格式为{task_group_id}.{定义operator时传入的原始task_id}，不再是代码里写的原始id。
• 分支判断函数update_pod_name_func里硬编码返回的是不带前缀的原始task_id（update_job_pod_name、skip_update_pod_name），BranchPythonOperator执行前会校验返回的task_id是否在当前DAG的任务列表中存在，找不到对应id的任务就会抛出该异常。

另外代码还存在一个隐性bug：SQL语句中混用了f-string和Jinja模板语法，f-string作用域内直接写{xcom_pull("pod_name")}会触发Python运行时错误，因为xcom_pull不是Python侧定义的函数，是Airflow Jinja模板渲染时才可用的方法，不能直接在f-string中调用。

内容的提问来源于stack exchange，提问作者WOWpopo