FastAPI多Gunicorn Worker启动时重复下载ML模型的解决方案咨询

Gunicorn多Worker模式下，每个Worker是独立进程，会同时执行启动逻辑，导致多个进程触发重复下载。无需引入外部存储，用本地文件系统的同步机制就能解决这个问题，以下是两种可行方案：

方案一：使用文件锁确保单一进程下载 #

基于Unix/Linux系统的fcntl模块实现跨进程排他锁，只有拿到锁的Worker会执行下载逻辑，其他Worker阻塞等待锁释放后再检查模型状态。

修改后的代码：

from contextlib import asynccontextmanager
from fastapi import FastAPI
import os
import fcntl

@asynccontextmanager
async def lifespan(app: FastAPI):
    model_path = os.path.join("./models", model)
    lock_path = "./models/model_download.lock"
    # 确保models目录存在
    os.makedirs("./models", exist_ok=True)
    
    # 初始化锁文件
    with open(lock_path, 'w'):
        pass
    
    with open(lock_path, 'r+') as lock_file:
        try:
            # 获取排他锁，其他进程会阻塞等待
            fcntl.flock(lock_file, fcntl.LOCK_EX)
            
            # 再次检查模型（防止等待期间其他进程已完成下载）
            if os.path.isdir(model_path):
                print(f"Startup: model {model} exists, skip download")
            else:
                print(f"Startup: install model {model}")
                download_model(model)
        finally:
            # 释放锁
            fcntl.flock(lock_file, fcntl.LOCK_UN)
    
    yield
    print('release models')

注：Windows系统需替换为msvcrt模块的锁机制，逻辑类似。

方案二：轮询等待模型（配合原子标记文件） #

通过原子操作创建下载标记文件，只有第一个进程能创建成功并执行下载，其他进程轮询等待模型出现或超时。

修改后的代码：

from contextlib import asynccontextmanager
from fastapi import FastAPI
import os
import time

@asynccontextmanager
async def lifespan(app: FastAPI):
    model_path = os.path.join("./models", model)
    downloading_flag = "./models/.downloading"
    os.makedirs("./models", exist_ok=True)
    
    if os.path.isdir(model_path):
        print(f"Startup: model {model} exists, skip download")
    else:
        try:
            # 原子创建标记文件，只有第一个进程能成功
            fd = os.open(downloading_flag, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.close(fd)
            # 执行下载
            print(f"Startup: install model {model}")
            download_model(model)
            # 下载完成删除标记
            os.remove(downloading_flag)
        except FileExistsError:
            # 其他进程进入轮询等待
            print(f"Startup: waiting for model {model} to be downloaded...")
            wait_time = 0
            max_wait = 300  # 最大等待5分钟，可按需调整
            while wait_time < max_wait:
                if os.path.isdir(model_path):
                    print(f"Startup: model {model} downloaded by another worker")
                    break
                time.sleep(5)
                wait_time +=5
            else:
                raise RuntimeError(f"Timeout waiting for model {model} to download")
    
    yield
    print('release models')

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