python - 从使用 Airflow KubernetesPodOperator 启动的 Pod 获取上下文

Question

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

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

谢谢!

Answer 1

我找到了一个很好的解决办法

我为 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)