Airflow 2.3.3中如何重试整个TaskGroup提交Spark K8s作业？

针对你遇到的问题——Spark作业执行失败（如UnknownHostException）时，仅重试传感器无法重新提交作业，需要触发整个TaskGroup的重试逻辑，以下是两种可行的实现方案：

方案1：通过回调函数重置提交任务状态 #

核心思路是当传感器监控失败时，清除同组内提交任务的成功状态，使其重新执行并提交新的Spark作业。

步骤1：定义重置提交任务的回调函数 #

添加以下代码到你的DAG文件中，用于清除指定提交任务的实例状态：

from airflow.api.common.experimental.clear import clear_task_instances
from airflow.models import TaskInstance
from airflow.utils.session import create_session

def reset_submit_task(context):
    ti = context["task_instance"]
    # 从传感器task_id推导提交任务的task_id（如将xxx-sensor转为xxx）
    submit_task_id = f"{ti.task.task_group_id}.{ti.task_id.replace('-sensor', '')}"
    dag_run = ti.dag_run

    with create_session() as session:
        submit_ti = session.query(TaskInstance).filter(
            TaskInstance.dag_id == ti.dag_id,
            TaskInstance.task_id == submit_task_id,
            TaskInstance.run_id == dag_run.run_id
        ).first()

        if submit_ti and submit_ti.state == "success":
            # 清除提交任务的状态，允许其重新执行
            clear_task_instances(
                task_ids=[submit_task_id],
                dag_id=ti.dag_id,
                run_id=dag_run.run_id,
                session=session,
                include_upstream=False,
                include_downstream=False,
                reset_dag_runs=False
            )

步骤2：修改传感器配置 #

移除传感器的retries=1配置，继承TaskGroup的重试策略，并绑定回调函数：

task_some_task_name_sensor = SparkKubernetesSensor(
    task_id='some-task-name-sensor',
    namespace="batch",
    application_name="{{ task_instance.xcom_pull(task_ids='tg-some-task-name.some-task-name')['metadata']['name'] }}",
    kubernetes_conn_id="kubernetes_default",
    dag=dag,
    attach_log=True,
    # 同时保留原有的Slack告警和新增的重置回调
    on_failure_callback=lambda ctx: [slack_alert(ctx), reset_submit_task(ctx)]
)

步骤3：确保作业名称唯一性 #

修改Spark作业的YAML文件，使用Airflow模板变量生成唯一名称，避免重复提交失败：

# k8s/some-task-name.yaml
apiVersion: sparkoperator.k8s.io/v1beta2
kind: SparkApplication
metadata:
  name: some-task-name-{{ run_id }}  # 用run_id保证名称唯一
  namespace: batch
spec:
  # 原有作业配置...

同时在提交任务中开启YAML模板渲染：

task_some_task_name = SparkKubernetesOperator(
    task_id='some-task-name',
    namespace="batch",
    application_file="k8s/some-task-name.yaml",
    do_xcom_push=True,
    dag=dag,
    template_ext=['.yaml']  # 启用YAML模板解析
)

方案2：用TaskFlow API封装提交+监控逻辑 #

将提交和监控作业的逻辑封装为单个TaskFlow任务，这样任务失败时会整体重试，自动重新提交新作业。

实现代码 #

from airflow.decorators import task

@task(
    task_id="submit_and_monitor",
    retries=5,
    retry_delay=timedelta(minutes=30),
    on_failure_callback=slack_alert
)
def submit_and_monitor_spark_job(application_file, namespace):
    from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator
    from airflow.providers.cncf.kubernetes.sensors.spark_kubernetes import SparkKubernetesSensor
    from airflow.models import TaskInstance, DagRun
    from airflow.utils.session import create_session

    # 初始化提交Operator
    submit_op = SparkKubernetesOperator(
        task_id="temp_submit",
        namespace=namespace,
        application_file=application_file,
        do_xcom_push=True,
        template_ext=['.yaml']
    )

    # 获取当前DAG运行实例
    with create_session() as session:
        dag_run = session.query(DagRun).filter(DagRun.dag_id == submit_op.dag_id).order_by(DagRun.execution_date.desc()).first()

    # 执行提交操作
    submit_result = submit_op.execute({
        "task_instance": TaskInstance(task=submit_op, run_id=dag_run.run_id),
        "dag_run": dag_run
    })
    app_name = submit_result["metadata"]["name"]

    # 初始化传感器并执行监控
    sensor_op = SparkKubernetesSensor(
        task_id="temp_sensor",
        namespace=namespace,
        application_name=app_name,
        kubernetes_conn_id="kubernetes_default",
        attach_log=True
    )
    sensor_op.execute({
        "task_instance": TaskInstance(task=sensor_op, run_id=dag_run.run_id),
        "dag_run": dag_run
    })

# 在TaskGroup中使用该任务
with TaskGroup("tg-some-task-name", default_args=default_args,) as tx_some_task_name:
    submit_and_monitor(
        application_file="k8s/some-task-name.yaml",
        namespace="batch"
    )

关键注意事项 #

• 必须保证每次提交的Spark作业名称唯一，否则K8s会因资源冲突拒绝新的提交请求。
• 方案1中使用Airflow API清除任务实例时，需确保Airflow运行账号拥有操作TaskInstance的权限。
• 两种方案均需移除传感器独立的重试配置，让重试逻辑统一由TaskGroup或封装后的任务控制。

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