Airflow PythonVirtualenvOperator虚拟环境加载DAG包内模块

当前报错的核心原因是两个执行上下文的模块搜索路径不互通：

• DAG定义阶段运行在Airflow主调度进程中，Airflow会自动将zip形式的DAG包临时解压到本地临时目录，并把解压根目录（也就是和dag.py同级、存放package/文件夹的目录）加入主进程的sys.path，因此直接导入package是正常的。
• PythonVirtualenvOperator执行任务时会创建全新的隔离虚拟环境进程，这个进程既不会继承Airflow主进程的路径配置，也默认不会把自定义业务模块所在目录纳入虚拟环境的模块搜索范围，因此在callable执行阶段导入package就会抛出模块不存在的错误。

方案1：使用Operator内置的路径注入能力（推荐，Airflow 2.0+适用） #

PythonVirtualenvOperator原生提供include_path参数，专门用于把指定路径列表直接注入新建虚拟环境的sys.path，不需要手动修改callable内的路径逻辑。
直接在DAG定义阶段动态获取自定义模块的根路径，传入参数即可：

import os
from datetime import datetime
from airflow import models
from airflow.operators.python_operator import PythonVirtualenvOperator

# 动态获取当前dag.py所在的绝对路径（即自定义package模块的同级根目录）
CUSTOM_MODULE_ROOT = os.path.dirname(os.path.abspath(__file__))

def my_function(**kwargs):
    # 路径已经被Operator自动注入，这里可以直接导入模块
    from package import my_module
    # 正常执行业务逻辑即可
    my_module.execute()

default_dag_args = {
    "start_date": datetime(2024, 1, 1)
}

with models.DAG(
    dag_id="packaged_virtualenv_dag",
    default_args=default_dag_args
) as dag:
    virtualenv_task = PythonVirtualenvOperator(
        task_id="virtualenv_python",
        python_callable=my_function,
        # 把自定义模块根路径注入虚拟环境的搜索路径
        include_path=[CUSTOM_MODULE_ROOT],
        system_site_packages=True,
        dag=dag,
    )

方案2：兼容老版本Airflow的手动路径注入 #

如果使用的是低于2.0版本的Airflow，没有include_path参数，可以通过任务参数把模块路径传递到虚拟环境进程内，在callable执行入口手动把路径加入sys.path：

import os
import sys
# 其他依赖导入逻辑同方案1

def my_function(**kwargs):
    # 第一行先把传入的自定义模块根路径加到虚拟环境的搜索路径首位
    sys.path.insert(0, kwargs["custom_module_root"])
    # 路径配置完成后再导入自定义模块
    from package import my_module
    my_module.execute()

# DAG定义部分同前，任务初始化时修改参数：
with models.DAG(
    # 配置同前
) as dag:
    virtualenv_task = PythonVirtualenvOperator(
        task_id="virtualenv_python",
        python_callable=my_function,
        # 通过任务参数把主进程拿到的模块路径传递给虚拟环境执行进程
        op_kwargs={"custom_module_root": CUSTOM_MODULE_ROOT},
        system_site_packages=True,
        dag=dag,
    )

• 不要用相对路径在虚拟环境内导入自定义模块：虚拟环境进程启动时的工作目录是系统随机生成的临时目录，不是DAG解压的根目录，相对路径会指向不存在的位置。
• 不要依赖system_site_packages=True访问自定义模块：这个参数仅作用于系统Python已安装的第三方包目录，你的自定义业务模块从未安装到系统site-packages路径下，开启该参数也无法解决模块找不到的问题。
• 不要尝试把自定义模块通过requirements参数走pip安装流程：该参数配置的依赖会从PyPI或私有源拉取安装，本地zip包内的业务模块无法被pip识别和安装，反而会增加不必要的流程复杂度。

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