Composer DAG导入google.cloud.aiplatform后无法并行执行任务

这既不是配额问题，也暂时不需要提交GCP支持工单，核心问题出在google.cloud.aiplatform库的导入时机和资源占用上：

为什么不是配额问题？ #

配额限制会触发明确的QuotaExceeded类报错，且你未导入该库时能正常并行50个任务，说明Airflow的并发配置（如concurrency=50）、底层执行器资源（如Celery Worker的并发数）是足够的，导入库后出现的并行失效和配额无关。

核心原因 #

google.cloud.aiplatform在导入阶段会执行全局初始化操作：加载认证凭据、初始化客户端实例、预加载部分资源。如果在DAG文件顶部全局导入该库，会导致：

• DAG解析阶段就加载了大量资源，占用Airflow调度器或Worker的内存；
• 若使用Celery等共享进程的执行器，多个任务共享进程空间时，库的全局初始化逻辑可能引发锁冲突，导致任务无法并行启动；
• 每个任务进程加载该库后内存占用飙升，Worker因资源不足被迫串行处理任务。

解决方案 #

1. 延迟导入库到任务函数内部
   不要在DAG文件顶部全局导入google.cloud.aiplatform，而是在PythonOperator的可调用函数内部导入，确保只有任务执行时才加载库，且每个任务进程独立初始化：

   def my_python_task(task_id):
       from google.cloud import aiplatform
       # 在这里编写你的Vertex AI业务逻辑
       print(f'Hello World {task_id}')
   

2. 检查Worker资源配置

   • 若使用Celery Executor，调整worker_concurrency参数，同时确保Worker节点有足够的CPU和内存资源，避免因单个任务内存占用过高导致无法并行；
   • 优先选择Kubernetes Executor，每个任务运行在独立Pod中，资源隔离更彻底，不会因单个库的资源占用影响其他任务并行。

3. 排查库版本问题
   尝试升级google-cloud-aiplatform到最新稳定版，部分旧版本可能存在初始化时的锁或资源泄漏问题。

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