Airflow升级至2.7.2后PythonVirtualEnvOperator DAG运行异常

问题背景 #

从Airflow 2.2升级到2.7.2后，原有PythonVirtualEnvOperator DAG出现两个连续错误：

1. 传递{{ data_interval_end }}到op_args时，触发错误：

   ValueError: The key 'data_interval_end' in args is a part of kwargs and therefore reserved.
   

2. 移除op_args及对应参数后，又出现序列化错误：

   file.write_bytes(self.pickling_library.dumps({"args": self.op_args, "kwargs": self.op_kwargs}))
   TypeError: cannot pickle 'module' object
   

原代码片段 #

def process(data_interval_end, dag_run):
     pass
    
        
    with DAG(
                    dag_id='da_survey_presentation',
                    default_args=DEFAULT_ARGS,
                    schedule_interval='0 8 * * 1',  # run once a week at midnight Pacific every Monday morning
                    max_active_runs=1,
                    catchup=False,
                    concurrency=1,
                    on_failure_callback=partial(on_failure, stage),
                    sla_miss_callback=partial(sla_miss, stage),
                    start_date=datetime.datetime(2022, 1, 1),
                    is_paused_upon_creation=(stage != 'prod'),
                    tags=['da']
             ) as dag:
                PythonVirtualenvOperator(
                    task_id="virtualenv_task",
                    python_callable=process,
                    op_args=[
                        '{{ data_interval_end }}'
                    ],
                    provide_context=True,
                    requirements=["google-api-python-client==1.6.7", "df2gspread==1.0.4"],
                    system_site_packages=True,
                    dag=dag
                )

解决步骤 #

1. 处理data_interval_end保留关键字问题 #

Airflow 2.x后期版本中，data_interval_end属于上下文自动注入的kwargs参数，不能通过op_args传递。需要改为显式通过op_kwargs传递模板化变量，同时调整函数参数接收方式。

2. 解决无法pickle模块对象问题 #

provide_context=True会自动注入大量上下文对象，其中部分对象（如模块实例）无法被序列化。因此要去掉provide_context=True，改为只传递实际需要的变量，避免注入不必要的复杂对象。

修改后的完整代码 #

from datetime import datetime
from functools import partial
from airflow import DAG
from airflow.operators.python import PythonVirtualenvOperator

# 假设以下变量已提前定义
DEFAULT_ARGS = {}
def on_failure(stage, context):
    pass
def sla_miss(stage, context):
    pass
stage = "dev"

def process(data_interval_end, dag_run_id):
    # 替换为你的业务逻辑
    print(f"数据区间结束时间: {data_interval_end}")
    print(f"DAG运行ID: {dag_run_id}")

with DAG(
    dag_id='da_survey_presentation',
    default_args=DEFAULT_ARGS,
    schedule_interval='0 8 * * 1',  # 每周一早上8点运行
    max_active_runs=1,
    catchup=False,
    concurrency=1,
    on_failure_callback=partial(on_failure, stage),
    sla_miss_callback=partial(sla_miss, stage),
    start_date=datetime(2022, 1, 1),
    is_paused_upon_creation=(stage != 'prod'),
    tags=['da']
) as dag:
    PythonVirtualenvOperator(
        task_id="virtualenv_task",
        python_callable=process,
        # 显式传递需要的模板化变量，避免注入复杂对象
        op_kwargs={
            "data_interval_end": "{{ data_interval_end }}",
            "dag_run_id": "{{ dag_run.run_id }}"
        },
        requirements=["google-api-python-client==1.6.7", "df2gspread==1.0.4"],
        system_site_packages=True,
        dag=dag
    )

额外说明 #

• 如果确实需要使用dag_run的其他属性，尽量只传递具体属性值（如run_id、execution_date），而不是整个dag_run对象，避免序列化失败。
• Airflow 2.x中provide_context=True已逐步被更灵活的显式变量传递替代，推荐使用op_kwargs结合模板语法获取上下文变量。

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