如何在Google Cloud Composer中导入自定义Python包

1. setup.py的核心作用 #

setup.py是Python包的标准构建配置文件，核心作用是把包含多个脚本的文件夹标记为可被Python解释器识别的合法安装包，你可以在里面定义包名、版本号、依赖的第三方库、需要打包的模块范围等信息，包安装完成后就不需要手动调整sys.path或者靠相对路径导入，Python可以全局识别这个包。 #

2. 自定义包的本地结构准备 #

先按照Python官方规范整理你的包结构，示例如下：

your_custom_package/
├── __init__.py  # 空文件即可，作用是标记当前文件夹为Python可识别的包
├── module1.py   # 你的业务脚本1
├── module2.py   # 你的业务脚本2
├── sub_package/ # 可选：子包目录
│   ├── __init__.py
│   └── sub_module.py
└── setup.py     # 包配置文件

最简可用的setup.py写法参考：

from setuptools import setup, find_packages

setup(
    name="your-org-custom-package", # 自定义你的组织包名
    version="0.1.0", # 包版本号
    packages=find_packages(), # 自动扫描所有子包
    # 如果有依赖第三方库可以加下面的配置
    # install_requires=["requests>=2.30.0", "pandas>=2.0.0"]
)

3. Cloud Storage 包结构配置 #

Cloud Composer绑定的GCS桶有固定的目录规则，你可以根据场景选两种配置方案：

• 方案1：直接上传到plugins目录（适合迭代快的测试、小团队场景）
  把整个整理好的your_custom_package文件夹直接上传到Composer对应GCS桶的/plugins根目录下即可，Composer默认会把plugins目录加入Python的sys.path，不需要额外执行安装操作。
• 方案2：打包为whl文件安装（适合正式生产环境）
  本地执行命令python setup.py bdist_wheel生成.whl格式的安装包，把生成的whl文件上传到GCS桶的/dags/dependencies目录，然后在Composer环境的PYPI包配置页，添加包的安装路径即可，路径格式为/home/airflow/gcsfuse/dags/dependencies/你的包文件名称.whl。

4. DAG中导入使用示例 #

两种配置方案的导入方式完全一致，直接按包名导入模块即可，示例如下：

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

# 导入自定义包的模块
from your_custom_package.module1 import your_business_func
from your_custom_package.sub_package.sub_module import sub_func

def call_custom_package(**kwargs):
    your_business_func()
    sub_func()
    return "自定义包调用成功"

with DAG(
    dag_id="test_custom_package",
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False
) as dag:
    test_task = PythonOperator(
        task_id="test_custom_package_call",
        python_callable=call_custom_package
    )

常见避坑说明 #

• 所有包和子包目录下必须添加__init__.py文件，否则Python不会将目录识别为可导入的包
• 用plugins目录方案时不要和Airflow自带包、公共第三方包重名，避免导入冲突
• 包更新后如果未生效，可以重启Composer的scheduler和worker节点触发配置重载

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