78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
from celery import Celery
|
|
from app.config import settings
|
|
|
|
from celery.signals import worker_process_init, worker_process_shutdown
|
|
|
|
from app.tasks.async_runner import run_async, close_loop
|
|
from app.models.base import engine
|
|
|
|
|
|
def _derive_redis_db(url: str, db_no: int) -> str:
|
|
if not url:
|
|
return url
|
|
import re
|
|
if re.search(r"/\d+$", url):
|
|
return re.sub(r"/\d+$", f"/{db_no}", url)
|
|
return url.rstrip("/") + f"/{db_no}"
|
|
|
|
|
|
broker_url = settings.CELERY_BROKER_URL or (_derive_redis_db(settings.REDIS_URL, 1) if settings.REDIS_URL else "")
|
|
backend_url = settings.CELERY_RESULT_BACKEND or (_derive_redis_db(settings.REDIS_URL, 2) if settings.REDIS_URL else "")
|
|
|
|
if broker_url:
|
|
celery_app = Celery("videogen")
|
|
celery_app.conf.update(
|
|
broker_url=broker_url,
|
|
result_backend=backend_url or broker_url,
|
|
task_serializer="json",
|
|
accept_content=["json"],
|
|
result_serializer="json",
|
|
timezone="Asia/Shanghai",
|
|
enable_utc=True,
|
|
task_soft_time_limit=600,
|
|
task_time_limit=900,
|
|
task_acks_late=True,
|
|
task_reject_on_worker_lost=True,
|
|
task_track_started=True,
|
|
worker_prefetch_multiplier=1,
|
|
broker_transport_options={
|
|
"visibility_timeout": 3600,
|
|
},
|
|
task_routes={
|
|
"generation.chatapi_create_generation_task": {"queue": "gen_chatapi_create"},
|
|
"generation.poll_generation_task": {"queue": "gen_provider_poll"},
|
|
"generation.download_generation_result_task": {"queue": "gen_result_download"},
|
|
"app.tasks.cleanup.*": {"queue": "default"},
|
|
},
|
|
)
|
|
celery_app.autodiscover_tasks(["app.tasks"])
|
|
else:
|
|
celery_app = None
|
|
|
|
|
|
@worker_process_init.connect
|
|
def on_worker_process_init(**kwargs):
|
|
"""
|
|
Linux prefork 子进程启动后执行。
|
|
|
|
目的：
|
|
1. 丢弃 fork 前可能继承的连接池状态。
|
|
2. 后续任务会在当前子进程自己的长期 event loop 上重新建连接池。
|
|
"""
|
|
try:
|
|
run_async(engine.dispose())
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@worker_process_shutdown.connect
|
|
def on_worker_process_shutdown(**kwargs):
|
|
"""
|
|
子进程退出前关闭连接池和 event loop。
|
|
"""
|
|
try:
|
|
run_async(engine.dispose())
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
close_loop() |