K8s部署Airflow：如何用KubernetesPodOperator并行处理挂载卷文件

在Kubernetes某命名空间部署的Airflow中，该命名空间挂载了存储卷，需要并行处理卷内特定目录下的所有文件，但受资源限制只能同时运行3个KubernetesPodOperator任务。当前核心困境：

• DAG无法直接访问挂载卷，无法提前获取文件名传递给Operator
• 如果让每个Operator自行读取目录下首个文件，会出现多个任务重复处理同一文件的问题

目标是实现这样的流水线：

1. download_sra任务完成后，启动3个并行的fasterq_dump任务
2. 每个任务处理一个文件，处理完成后自动取下一个未处理的文件继续执行
3. 示例场景：目录有4个文件f1.sra/f2.sra/f3.sra/f4.sra，初始3个任务分别处理前3个，某任务完成后立即处理剩余的f4.sra

附当前KubernetesPodOperator示例代码：

fasterq_dump_instance_1 = KubernetesPodOperator(
        name="fasterq_dump_instance_1",
        task_id="fasterq_dump_instance_1",
        namespace="namespace-name",
        in_cluster=True,
        is_delete_operator_pod=True,
        get_logs=True,
        image="imagename",
        image_pull_policy="IfNotPresent",
        resources=k8s.V1ResourceRequirements(
            limits={"cpu": "2", "memory": "10Gi"},
            requests={"cpu": "1", "memory": "1Gi"},
        ),

        cmds=["python", "fasterq_dump.py", filename_to_process],

        volume_mounts=[
            V1VolumeMount(mount_path="/uftp", name="uftp-namespace"),
        ],
        volumes=[
            V1Volume(
                name="uftp-namespace",
                persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(
                    claim_name="uftp-namespace"
                )
            ),
        ],
    )

方案一：基于文件锁的原子性文件获取 #

修改fasterq_dump.py脚本，实现原子性获取未处理文件的逻辑，从根源避免多任务抢文件：

• 在挂载卷目录下创建lock子目录，用于存放临时锁文件
• 脚本启动后遍历目标目录下的.sra文件，对每个文件尝试用原子性方式创建锁文件（依赖os.open的O_CREAT | O_EXCL参数，保证同一时间只有一个任务能创建成功）
• 成功拿到锁的任务处理对应文件，完成后删除原文件和锁文件，再重新遍历目录取下一个未处理文件，直到无文件可处理

示例脚本核心逻辑：

import os
import time

TARGET_DIR = "/uftp"
LOCK_DIR = os.path.join(TARGET_DIR, "locks")
os.makedirs(LOCK_DIR, exist_ok=True)

def get_next_file():
    for filename in os.listdir(TARGET_DIR):
        if filename.endswith(".sra"):
            lock_path = os.path.join(LOCK_DIR, f"{filename}.lock")
            # 原子性创建锁文件，避免多任务争抢
            try:
                fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_RDWR)
                os.close(fd)
                return filename
            except OSError:
                # 锁文件已存在，跳过该文件
                continue
    return None

# 循环处理文件直到目录为空
while True:
    file_to_process = get_next_file()
    if not file_to_process:
        break
    # 执行你的fasterq_dump处理逻辑
    process_file(os.path.join(TARGET_DIR, file_to_process))
    # 处理完成后清理文件和锁
    os.remove(os.path.join(TARGET_DIR, file_to_process))
    os.remove(os.path.join(LOCK_DIR, f"{file_to_process}.lock"))

修改KubernetesPodOperator的cmds，不需要传递文件名，直接执行脚本：

cmds=["python", "fasterq_dump.py"]

最后在DAG中定义3个完全相同的任务，设置与download_sra的依赖：

download_sra = ... # 你的download任务

fasterq_dump_tasks = []
for i in range(3):
    task = KubernetesPodOperator(
        name=f"fasterq_dump_instance_{i+1}",
        task_id=f"fasterq_dump_instance_{i+1}",
        namespace="namespace-name",
        in_cluster=True,
        is_delete_operator_pod=True,
        get_logs=True,
        image="imagename",
        image_pull_policy="IfNotPresent",
        resources=k8s.V1ResourceRequirements(
            limits={"cpu": "2", "memory": "10Gi"},
            requests={"cpu": "1", "memory": "1Gi"},
        ),
        cmds=["python", "fasterq_dump.py"],
        volume_mounts=[
            V1VolumeMount(mount_path="/uftp", name="uftp-namespace"),
        ],
        volumes=[
            V1Volume(
                name="uftp-namespace",
                persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(
                    claim_name="uftp-namespace"
                )
            ),
        ],
    )
    download_sra >> task
    fasterq_dump_tasks.append(task)

方案二：使用共享任务队列 #

如果需要更清晰的任务调度逻辑，可以在挂载卷中维护一个任务队列文件：

1. download_sra任务完成后，生成一个task_queue.txt文件，每行写入一个待处理的文件名
2. 每个fasterq_dump任务启动后，通过文件锁原子性地从队列头部取出一个文件名
3. 处理完成后重复取队列中的文件，直到队列为空

示例队列处理逻辑：

import os
import fcntl

QUEUE_FILE = "/uftp/task_queue.txt"

def get_next_from_queue():
    with open(QUEUE_FILE, "r+") as f:
        # 加独占锁保证原子性操作
        fcntl.flock(f, fcntl.LOCK_EX)
        lines = f.readlines()
        if not lines:
            fcntl.flock(f, fcntl.LOCK_UN)
            return None
        # 取出队列第一个文件名
        filename = lines[0].strip()
        # 更新队列文件，移除已取走的行
        f.seek(0)
        f.writelines(lines[1:])
        f.truncate()
        fcntl.flock(f, fcntl.LOCK_UN)
    return filename

while True:
    file_to_process = get_next_from_queue()
    if not file_to_process:
        break
    process_file(os.path.join("/uftp", file_to_process))
    os.remove(os.path.join("/uftp", file_to_process))

需要在download_sra的脚本末尾添加生成队列的逻辑：

import os
TARGET_DIR = "/uftp"
with open("/uftp/task_queue.txt", "w") as f:
    for filename in os.listdir(TARGET_DIR):
        if filename.endswith(".sra"):
            f.write(f"{filename}\n")

方案三：Airflow动态任务生成（可选） #

如果文件数量固定，也可以先通过临时Pod扫描挂载卷获取文件名，再动态生成对应任务并限制并行度：

1. 用PythonOperator启动一个临时Kubernetes Pod，扫描挂载卷获取所有待处理文件名
2. 基于文件名列表动态生成KubernetesPodOperator任务，通过TaskGroup设置并行度为3

这种方式的缺点是无法处理后续新增的文件，仅适合文件数量固定的场景。

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