在Gunicorn部署中PyTorch DataLoader抛出多进程异常

核心原因 #

Gunicorn默认以fork模式启动工作进程，PyTorch DataLoader的多进程（num_workers>0）在fork出的子进程中会继承父进程未关闭的文件描述符，触发IO阻塞——这一问题在SageMaker的容器受限环境中会被放大，而Jupyter Notebook是单进程运行环境，因此不会出现该异常。

可行的代码调整方案 #

1. 在DataLoader初始化前关闭冗余文件描述符 #

手动关闭父进程中不必要的文件描述符，避免子进程继承后引发冲突：

import os
import resource

def close_unused_fds():
    max_fd = resource.getrlimit(resource.RLIMIT_NOFILE)[0]
    # 从fd=3开始关闭（0/1/2对应标准输入/输出/错误）
    for fd in range(3, max_fd):
        try:
            os.close(fd)
        except OSError:
            pass

# 创建DataLoader前执行
close_unused_fds()
data_loader = torch.utils.data.DataLoader(dataset, num_workers=2, batch_size=32)

2. 强制使用forkserver多进程启动方式 #

相比spawn，forkserver能避免子进程继承父进程的大部分资源，更适配Gunicorn的fork模式：

import torch.multiprocessing as mp

# 在模型/数据加载前设置启动方式
mp.set_start_method('forkserver', force=True)
# 限制每个子进程的线程数，避免资源争抢
torch.set_num_threads(1)
torch.set_num_interop_threads(1)

3. 自定义DataLoader的worker初始化函数 #

在每个worker启动时重置环境并清理资源，进一步隔离进程间的干扰：

def worker_init_fn(worker_id):
    # 重置随机种子，保证多进程推理一致性
    import numpy as np
    import random
    seed = torch.initial_seed() % 2**32
    np.random.seed(seed)
    random.seed(seed)
    # 关闭当前worker的冗余文件描述符
    close_unused_fds()  # 复用之前定义的close_unused_fds函数

# 初始化DataLoader时指定该函数
data_loader = torch.utils.data.DataLoader(
    dataset,
    num_workers=2,
    worker_init_fn=worker_init_fn,
    batch_size=32
)

4. 延迟初始化DataLoader #

避免在Gunicorn主进程中预加载模型或数据集，将DataLoader的初始化逻辑放在SageMaker的model_fn或predict_fn中——确保代码在Gunicorn worker进程启动后再执行，而非脚本全局初始化阶段。

验证建议 #

• 优先尝试forkserver启动方式+线程数限制，这是最直接的调整方案
• 若仍报错，叠加worker_init_fn关闭文件描述符的逻辑
• 测试时从num_workers=1开始逐步增加，观察稳定性

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