如何在Airflow任务间复用重型NLP服务实例？

针对你在Airflow分布式任务中复用NLP模型的需求，以下是几种可行的解决思路：

1. Worker进程级初始化加载模型 #

Airflow的Worker进程会持续处理多个任务，我们可以在Worker启动时一次性加载模型，让该进程内的所有任务复用同一个实例：

• 编写一个模型初始化模块，在Worker启动时执行初始化：

# nlp_service.py
nlp_service = None

def init_nlp_service():
    global nlp_service
    if not nlp_service:
        from your_module import SentenceService
        nlp_service = SentenceService(model="abc")

• 在Airflow配置文件airflow.cfg中指定Worker初始化脚本：

worker_initialization_script = /path/to/your_init_script.py

• 初始化脚本your_init_script.py中调用初始化函数：

from nlp_service import init_nlp_service
init_nlp_service()

• 任务函数中直接复用已初始化的实例：

from nlp_service import nlp_service

def process_text(text):
    doc = nlp_service.nlp(text)
    # 处理逻辑
    return processed_result

每个Worker进程启动时仅加载一次模型，该进程内的所有.expand任务都会复用这个实例，既保留了分布式能力，又避免重复加载开销。

2. 进程内单例/缓存实现复用 #

利用Python的进程内缓存机制，让每个Worker进程内仅创建一次模型实例，无需修改Airflow配置：

• 修改服务类，通过lru_cache实现进程内单例：

from functools import lru_cache
import spacy

class SentenceService:
    def __init__(self, model: str = "abc"):
        self.nlp = spacy.load(model)

@lru_cache(maxsize=1)
def get_sentence_service(model: str = "abc"):
    return SentenceService(model)

• 任务函数中调用缓存的服务实例：

def process_text(text):
    service = get_sentence_service()
    doc = service.nlp(text)
    # 处理逻辑
    return processed_result

lru_cache的作用域是单个进程，因此每个Worker进程首次调用时加载模型，后续任务直接复用已创建的实例。

3. 部署独立模型服务 #

将NLP模型部署为独立的API服务（如FastAPI、Flask），Airflow任务通过网络请求调用模型，实现跨Worker的全局复用：

• 编写模型API服务（示例用FastAPI）：

from fastapi import FastAPI
import spacy

app = FastAPI()
nlp = spacy.load("abc")

@app.post("/process-text")
def process_text(text: str):
    doc = nlp(text)
    # 处理逻辑
    return {"result": your_processed_result}

• Airflow任务中调用API：

import requests

def process_text(text):
    response = requests.post("http://your-model-service:8000/process-text", json={"text": text})
    return response.json()["result"]

此方案适合模型体积极大、Worker数量众多的场景，模型仅需加载一次，但会引入网络调用开销。

补充说明 #

你之前尝试的全局变量、单例失效的核心原因是：Airflow的分布式任务会运行在不同的Worker进程（甚至不同机器）中，进程间内存隔离，跨进程无法共享实例。上述方案要么在单个进程内复用，要么借助外部服务实现全局复用，均能解决你的问题。

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