如何在Airflow中用KubernetesPodOperator实现Kubernetes Job并行执行

完全可以实现，不需要创建N个KubernetesPodOperator，而是利用Airflow默认的算子（比如BashOperator或PythonOperator）直接创建Kubernetes Job资源，复用你熟悉的parallelism和completions参数来批量启动Pod。

方案1：使用BashOperator执行kubectl命令 #

如果Airflow的执行环境（Worker节点或Pod）已经安装了kubectl并配置了集群访问权限，这是最直接的方式。你可以通过环境变量传递动态的并行数，结合envsubst替换YAML模板中的参数，然后创建Job。

示例代码：

from airflow import DAG
from airflow.operators.bash import BashOperator
from datetime import datetime

with DAG(
    dag_id="k8s_batch_job",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    params={"parallel_num": 100000},  # 动态传入并行数
) as dag:
    create_k8s_job = BashOperator(
        task_id="create_batch_job",
        bash_command="""
        # 定义Job的YAML模板，替换参数
        cat << EOF | envsubst | kubectl apply -f -
apiVersion: batch/v1
kind: Job
metadata:
  annotations:
    eks.amazonaws.com/role-arn: "arn:aws:iam::{number}:role/{access_name}"
  name: test-job-{{ ts_nodash }}  # 用时间戳避免重名
  namespace: analytics
spec:
  completions: $PARALLEL_NUM
  parallelism: $PARALLEL_NUM
  ttlSecondsAfterFinished: 3600  # 任务完成后自动清理Job
  template:
    spec:
      containers:
      - name: your-container
        image: your-image:tag
        command: ["your-command"]
      restartPolicy: OnFailure
EOF
        """,
        env={"PARALLEL_NUM": "{{ params.parallel_num }}"},
    )

方案2：使用PythonOperator动态生成Job配置 #

如果需要更灵活的参数处理（比如动态生成不同的Job名称、调整其他配置），可以用PythonOperator结合subprocess模块执行kubectl命令：

示例代码：

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
import subprocess
import yaml

def create_k8s_job(**context):
    parallel_num = context["params"]["parallel_num"]
    job_config = {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "metadata": {
            "annotations": {
                "eks.amazonaws.com/role-arn": "arn:aws:iam::{number}:role/{access_name}"
            },
            "name": f"test-job-{context['ts_nodash']}",
            "namespace": "analytics"
        },
        "spec": {
            "completions": parallel_num,
            "parallelism": parallel_num,
            "ttlSecondsAfterFinished": 3600,
            "template": {
                "spec": {
                    "containers": [
                        {
                            "name": "your-container",
                            "image": "your-image:tag",
                            "command": ["your-command"]
                        }
                    ],
                    "restartPolicy": "OnFailure"
                }
            }
        }
    }
    # 将字典转为YAML字符串，通过kubectl创建
    yaml_str = yaml.dump(job_config)
    subprocess.run(
        ["kubectl", "apply", "-f", "-"],
        input=yaml_str.encode("utf-8"),
        check=True
    )

with DAG(
    dag_id="k8s_batch_job_python",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    params={"parallel_num": 100000},
) as dag:
    create_job_task = PythonOperator(
        task_id="create_batch_job",
        python_callable=create_k8s_job,
        provide_context=True
    )

关键注意事项 #

• 确保Airflow执行任务的环境拥有kubectl工具，并且具备在analytics namespace创建Job的权限（可以通过ServiceAccount绑定RBAC角色实现）。
• 用ts_nodash作为Job名称后缀，避免重复创建同名Job导致冲突。
• 设置ttlSecondsAfterFinished参数可以自动清理完成后的Job，避免集群资源浪费。
• 动态参数可以通过Airflow的Params、Variables或DAG运行时配置传入，适配不同的并行数需求。

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