26 lines
760 B
Python
26 lines
760 B
Python
import atexit
|
|
import logging
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 全局初始化线程池
|
|
_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix='image_embedding_')
|
|
|
|
def run_embedding_task(fn, *args, **kwargs):
|
|
"""
|
|
提交后台任务到线程池
|
|
"""
|
|
logger.info("Submitting embedding task to background thread...")
|
|
return _executor.submit(fn, *args, **kwargs)
|
|
|
|
def _shutdown_executor():
|
|
"""
|
|
优雅关闭线程池，在 Gunicorn worker 退出时触发
|
|
"""
|
|
logger.info("Shutting down background thread pool...")
|
|
_executor.shutdown(wait=False)
|
|
|
|
# 注册到系统退出事件，这样就不会报 _shutdown not defined 了
|
|
atexit.register(_shutdown_executor)
|