首页 > 解决方案 > 从使用 Airflow KubernetesPodOperator 启动的 Pod 获取上下文

问题描述

我们有一些使用 KubernetesPodOperator 启动 pod 的 dag,我正在尝试获取 pod 内部的一些信息,例如 dag_id、task_id、try_number、环境等。

我知道我可以从 Airflow 任务的上下文中获取此信息(例如,Python Operator 上的 kwargs),但我一直在想,有没有办法可以从启动的 pod 中获取该上下文?

谢谢!

标签: pythonkubernetesairflow

解决方案


我找到了一个很好的解决方案

我为类 KubernetesPodOperator 制作了一个自定义包装器,并使用 Airflow 任务的上下文更新了 env_vars

import airflow.configuration as config
from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator as AirflowKubernetesPodOperator

class KubernetesPodOperator(AirflowKubernetesPodOperator):

    def execute(self, context):

        environment = config.conf.get('webserver', 'web_server_name')

        ti = context['ti']
        dag_id = ti.dag_id
        task_id = ti.task_id
        run_id = context['run_id']
        try_number = str(ti._try_number)

        labels = { 
            'ENVIRONMENT' : environment,
            'DAG_ID'      : dag_id, 
            'TASK_ID'     : task_id, 
            'RUN_ID'      : run_id,
            'TRY_NUMBER'  : try_number,
        }

        self.env_vars.update(labels)
        super().execute(context)

推荐阅读