PlantAI-dev/celery_app.py

import os
import sys
from pathlib import Path

from celery import Celery
from kombu import Exchange, Queue

LOCAL_PROJECT_ROOT = Path(__file__).resolve().parent
if str(LOCAL_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(LOCAL_PROJECT_ROOT))

from gfa_core.project_paths import PROJECT_ROOT

os.environ.setdefault("PROJECT_ROOT", str(PROJECT_ROOT))
os.environ.setdefault("BBAV5_PROJECT_ROOT", str(PROJECT_ROOT))
QUEUE_PREFIX = os.environ.get("BBAV5_QUEUE_PREFIX", "plantai_dev_9000").strip() or "plantai_dev_9000"
REDIS_URL = os.environ.get("BBAV5_REDIS_URL", "redis://localhost:6379/10")


def q(name: str) -> str:
    return f"{QUEUE_PREFIX}.{name}"


TASK_EXCHANGE = Exchange(q("tasks"), type="direct")


def queue(name: str) -> Queue:
    queue_name = q(name)
    return Queue(queue_name, TASK_EXCHANGE, routing_key=queue_name)

app = Celery(
    "plantai_dev_9000",
    broker=REDIS_URL,
    backend=REDIS_URL,
    include=["tasks"],
)

app.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    timezone="Asia/Shanghai",
    enable_utc=True,
    task_track_started=True,
    task_default_queue=q("default"),
    task_default_exchange=q("tasks"),
    task_default_exchange_type="direct",
    task_default_routing_key=q("default"),
    task_queues=(
        queue("default"),
        queue("heavy"),
        queue("medium"),
        queue("light"),
        queue("rnaseq"),
    ),
    task_routes={
        "tasks.execute_agent_task": {"queue": q("default"), "routing_key": q("default")},
    },
    task_time_limit=172800,
    task_soft_time_limit=169200,
    task_annotations={
        "tasks.execute_agent_task": {"time_limit": 3600, "soft_time_limit": 3300},
        "tasks.execute_single_stage_task": {"time_limit": 172800, "soft_time_limit": 169200},
        "tasks.execute_lightweight_extension_operation": {
            "time_limit": None,
            "soft_time_limit": None,
        },
    },
    worker_prefetch_multiplier=1,
)

if __name__ == "__main__":
    app.start()