Airflow运行CUDA和Torch报错：无法在fork子进程初始化CUDA

问题说明 #

运行Airflow DAG时触发CUDA初始化报错：

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

核心原因是Gunicorn默认采用fork模式管理进程，而Torch的CUDA操作要求使用spawn启动多进程，两者冲突导致无法正常运行GPU任务。

DAG代码示例 #

from airflow import DAG
from airflow.decorators import task
from datetime import datetime

with DAG(
    dag_id='dag_temp',
    description='Dag temp',
    start_date=datetime(2022, 8, 18),
    catchup=False
) as dag:

    from transformers import AutoTokenizer, RobertaModel
    import torch
    
    @task()
    def load_models():
        device = 'cuda' if torch.cuda.is_available() else 'cpu' # device = cuda

        tokenizer = AutoTokenizer.from_pretrained("roberta-base")
        model = RobertaModel.from_pretrained("roberta-base")
        model.to(device)

完整报错栈 #

Traceback (most recent call last):
  File "/home/leobit/Development/mlfactory-pipeline/src/dags/med_reader_pipeline.py", line 34, in load_models
    model.to(device)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1749, in to
    return super().to(*args, **kwargs)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 989, in to
    return self._apply(convert)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  [Previous line repeated 1 more time]
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 664, in _apply
    param_applied = fn(param)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 987, in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
  File "/home/leobit/.cache/pypoetry/virtualenvs/mlfactory-pipeline-LOOcmODM-py3.10/lib/python3.10/site-packages/torch/cuda/__init__.py", line 217, in _lazy_init
    raise RuntimeError(
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

已尝试的方案 #

• 将execute_tasks_new_python_interpreter配置改为True
• 减少Gunicorn的worker数量
• 设置torch.multiprocessing为spawn启动方法：mp.set_start_method('spawn', force=True)

可行解决方案 #

1. 调整Airflow任务执行模式 #

确保Airflow使用支持spawn的执行器：

• 若使用LocalExecutor，在Airflow配置文件中添加：

  [core]
  executor = LocalExecutor
  [local_executor]
  parallelism = 4
  

• 关键原则：所有CUDA相关的导入、初始化操作必须放在任务函数内部，禁止在DAG全局作用域中提前导入Torch或初始化CUDA。

2. 使用KubernetesExecutor（推荐） #

将每个任务部署为独立K8s Pod，彻底隔离进程环境，从根源避免fork带来的CUDA冲突。每个Pod可单独申请GPU资源，任务内直接正常使用Torch CUDA即可。

3. 修改Gunicorn启动参数 #

强制Gunicorn使用非fork的worker模式，例如使用线程型worker：

gunicorn --worker-class=gthread --workers=2 --threads=4 airflow.webserver:app

或使用gevent worker类（需确保Torch与gevent兼容）。

4. 优化任务代码结构 #

调整代码，将CUDA相关逻辑完全封装在任务函数内部：

from airflow import DAG
from airflow.decorators import task
from datetime import datetime

with DAG(
    dag_id='dag_temp',
    description='Dag temp',
    start_date=datetime(2022, 8, 18),
    catchup=False
) as dag:

    @task()
    def load_models():
        # 所有CUDA相关操作放在任务内部
        import torch
        import multiprocessing as mp
        from transformers import AutoTokenizer, RobertaModel
        
        # 提前设置spawn启动方法
        mp.set_start_method('spawn', force=True)
        
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        tokenizer = AutoTokenizer.from_pretrained("roberta-base")
        model = RobertaModel.from_pretrained("roberta-base")
        model.to(device)
        # 后续业务逻辑

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