如何从失败的Airflow KubernetesPodOperator任务拉取XCOM退出码

核心原因 #

你之前的方法无效，是因为KPO默认仅在任务成功执行完成时才会触发XCOM推送逻辑——失败时execute方法直接抛出异常，根本没走到推送XCOM的代码段。即使事后通过回调把任务状态改成成功，也无法补触发原本的XCOM推送流程。

可行解决方案 #

1. 自定义算子时强制推送XCOM（推荐）

继承KPO时，重写execute方法，在任务执行完成（无论成功失败）后主动推送退出码到XCOM：

from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

class ErrorHandlingKPO(KubernetesPodOperator):
    def execute(self, context):
        exit_code = 0
        try:
            # 执行父类的核心逻辑
            super().execute(context)
        except Exception as e:
            # 从KPO抛出的异常中提取退出码（KPO失败时会携带exit_code属性）
            exit_code = getattr(e, 'exit_code', -1)
            # 保留原有的失败逻辑，不要在这里捕获异常（否则任务会被标记为成功）
            raise
        finally:
            # 无论成功失败，强制推送退出码到XCOM
            self.xcom_push(context, key='task_exit_code', value=exit_code)

这样后续任务就能通过ti.xcom_pull(task_ids='your_kpo_task_id', key='task_exit_code')获取到退出码，不管原任务是成功还是失败。

2. 通过失败回调推送XCOM

如果不想修改算子代码，可以给KPO任务添加on_failure_callback，在回调中提取退出码并推送：

def push_exit_code_callback(context):
    task_instance = context['task_instance']
    # 从失败异常中提取退出码
    exit_code = getattr(task_instance.exception, 'exit_code', -1)
    # 推送XCOM
    task_instance.xcom_push(key='task_exit_code', value=exit_code)

# 定义KPO任务
test_task = KubernetesPodOperator(
    task_id='test_error_handling',
    image='your-image:latest',
    cmds=['bash', '-c', 'exit 1'],  # 模拟失败场景
    on_failure_callback=push_exit_code_callback,
    # 其他必要参数...
)

3. 调整任务状态的同时补推XCOM

如果一定要用“把失败任务标记为成功”的思路，需要在回调中同时完成状态修改和XCOM推送：

def mark_success_and_push_exit_code(context):
    task_instance = context['task_instance']
    # 提取退出码
    exit_code = getattr(task_instance.exception, 'exit_code', -1)
    # 推送XCOM
    task_instance.xcom_push(key='task_exit_code', value=exit_code)
    # 修改任务状态为成功
    task_instance.set_state('success')

test_task = KubernetesPodOperator(
    task_id='test_task',
    # 其他参数...
    on_failure_callback=mark_success_and_push_exit_code
)

这个方法需要注意：修改状态后，Airflow会认为任务成功，可能会影响DAG的整体调度逻辑，仅适合测试场景使用。

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