如何限制Airflow任务资源，或阻止在节点本地运行Spark作业的任务？

一、阻止本地模式Spark任务启动 #

1. 自定义SparkSubmitOperator拦截 #

编写Airflow插件，替换原生SparkSubmitOperator，在任务执行前校验Spark运行模式，直接禁止local模式启动：

from airflow.plugins_manager import AirflowPlugin
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
from airflow.utils.decorators import apply_defaults

class RestrictedSparkSubmitOperator(SparkSubmitOperator):
    @apply_defaults
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def execute(self, context):
        # 校验直接传入的master参数
        if self.master and self.master.startswith('local'):
            raise ValueError("禁止在Airflow节点以本地模式运行Spark作业，请使用集群模式")
        # 校验conf配置中的master参数
        if self.conf and 'spark.master' in self.conf and self.conf['spark.master'].startswith('local'):
            raise ValueError("禁止在Airflow节点以本地模式运行Spark作业，请使用集群模式")
        super().execute(context)

class SparkRestrictionPlugin(AirflowPlugin):
    name = "spark_restriction_plugin"
    operators = [RestrictedSparkSubmitOperator]

要求所有团队使用这个自定义Operator提交Spark作业，从任务提交层面拦截违规操作。

2. 全局任务钩子拦截 #

利用Airflow的任务运行钩子，在任务启动前强制检查并阻止：

from airflow.plugins_manager import AirflowPlugin
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator

def block_local_spark(context):
    task_instance = context['task_instance']
    task = task_instance.task
    if isinstance(task, SparkSubmitOperator):
        if (task.master and task.master.startswith('local')) or \
           (task.conf and 'spark.master' in task.conf and task.conf['spark.master'].startswith('local')):
            task_instance.set_state('failed')
            raise ValueError("本地模式Spark作业已被系统拦截")

class SparkHookPlugin(AirflowPlugin):
    name = "spark_hook_plugin"
    task_instance_running = [block_local_spark]

3. Spark节点配置强制覆盖 #

在Airflow节点的spark-defaults.conf中添加全局配置：

spark.master spark://your-spark-cluster-master:7077

该配置会强制覆盖任务指定的local模式，从Spark层面限制本地运行。

二、禁止违规DAG导入 #

1. DAG导入校验插件 #

编写插件在DAG加载阶段检查配置，发现违规任务直接阻止DAG导入：

from airflow.plugins_manager import AirflowPlugin
from airflow.models import DAG
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator

def validate_dag_config(dag: DAG):
    for task in dag.tasks:
        if isinstance(task, SparkSubmitOperator):
            if (task.master and task.master.startswith('local')) or \
               (task.conf and 'spark.master' in task.conf and task.conf['spark.master'].startswith('local')):
                raise ImportError(f"DAG [{dag.dag_id}] 包含禁止的本地模式Spark任务，无法完成导入")

class DAGValidationPlugin(AirflowPlugin):
    name = "dag_validation_plugin"
    dag_import_hooks = [validate_dag_config]

2. CI/CD前置校验 #

在DAG代码仓库中添加CI/CD流程（如GitHub Actions、GitLab CI），在代码合并前自动扫描DAG文件，检查是否存在本地模式的SparkSubmitOperator配置，发现则直接阻止合并。

三、限制Airflow任务资源 #

1. CeleryExecutor资源限制 #

如果使用CeleryExecutor，启动Worker时添加资源限制参数：

airflow celery worker --concurrency=2 --max-memory-per-child=8192 --max-tasks-per-child=10

• --concurrency：限制Worker同时处理的任务数量
• --max-memory-per-child：限制单个任务进程的最大内存（单位：MB）
• --max-tasks-per-child：限制Worker进程处理的任务数，防止内存泄漏堆积

2. KubernetesExecutor资源隔离 #

如果基于Kubernetes部署Airflow，通过Pod模板配置资源请求与限制：

# pod_template.yaml
apiVersion: v1
kind: Pod
metadata:
  name: airflow-task-pod
spec:
  containers:
    - name: airflow-worker
      image: your-airflow-image:latest
      resources:
        requests:
          cpu: "1"
          memory: "4Gi"
        limits:
          cpu: "4"
          memory: "16Gi"

在airflow.cfg中指定模板路径：

[kubernetes]
pod_template_file = /path/to/pod_template.yaml

每个任务运行在独立Pod中，资源限制不会影响Airflow核心节点。

3. Airflow核心配置优化 #

在airflow.cfg中添加以下配置：

[core]
worker_max_tasks_per_child = 10  # 每个Worker进程处理10个任务后自动重启
dag_run_timeout = 3600  # DAG最长运行时间（单位：秒）

[scheduler]
max_tis_per_query = 50  # 调度器单次查询的任务实例数，避免调度过载

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