""" app/workers/tasks.py — Celery 任务入口(编排层) 只接收 task_id(基石B)。具体步骤在 pipeline_steps / pipeline_io。 打包任务在 packaging_task.py,此处 re-export 保持旧引用不变。 """ import logging from app.workers.celery_app import celery_app # re-export 保持旧代码 `from app.workers.tasks import build_delivery_package` 不变 from app.workers.packaging_task import build_delivery_package # noqa: F401 from app.workers.replenish_task import replenish_text_candidates # noqa: F401 logger = logging.getLogger(__name__) def _json_loads_safe(raw: str | None) -> dict: """TextCandidate.content 存整条 JSON dict(title/content/tags/...)。 解析失败兼容旧数据:纯正文则包成 {'content': raw}。供 R2 重生复用已有文案。""" import json as _json if not raw: return {} try: v = _json.loads(raw) return v if isinstance(v, dict) else {"content": raw} except (ValueError, TypeError): return {"content": raw} def _get_db(): """获取同步 DB Session(Celery worker 环境用同步)""" from app.core.database import SessionLocal return SessionLocal() def _push_event_sync(task_id: int, workspace_id: int, event: str, data: dict, seq: int): """同步推 SSE 事件(worker 内部用)""" import json as _json import redis as redis_lib from app.core.config import get_settings s = get_settings() r = redis_lib.from_url(s.REDIS_URL, decode_responses=True) hist_key = f"sse:task:{task_id}:events" channel = f"sse:task:{task_id}" record = _json.dumps({"event": event, "data": data, "seq": seq, "workspace_id": workspace_id}, ensure_ascii=False) r.rpush(hist_key, record) r.expire(hist_key, 3600) r.publish(channel, record) def _acquire_run_lock(task_id: int) -> bool: """ 幂等锁:同一 task 同时只允许一条 pipeline 在跑。 visibility_timeout 是主防线(防 broker 误判重投);此锁是双保险, 挡住"窗口内仍并发重投"导致两个 worker 同跑同一 task、重复烧钱(task75 教训)。 返回 True=抢到锁可执行;False=已有人在跑,本次直接跳过。 锁 TTL 略大于 visibility_timeout,确保跨整个任务生命周期。 """ import redis as redis_lib from app.core.config import get_settings r = redis_lib.from_url(get_settings().REDIS_URL, decode_responses=True) # NX+EX:不存在才设,自带过期防死锁(worker 崩了锁自动释放) return bool(r.set(f"pipeline:lock:{task_id}", "1", nx=True, ex=7500)) def _release_run_lock(task_id: int) -> None: import redis as redis_lib from app.core.config import get_settings r = redis_lib.from_url(get_settings().REDIS_URL, decode_responses=True) r.delete(f"pipeline:lock:{task_id}") @celery_app.task( bind=True, name="app.workers.tasks.run_generation_pipeline", max_retries=3, default_retry_delay=30, queue="generation", ) def run_generation_pipeline(self, task_id: int, regen_strategy: str | None = None, regen_role: str | None = None, custom_prompt: str | None = None) -> dict: """ 生产链主任务。只接收 task_id,绝不接收 key。 子步骤委托给 pipeline_steps(查库/解密/飞轮上下文) 和 pipeline_io(文案/图片生成+存储)。 R2 重生参数(均 None=常规全量生成): regen_strategy='A'/'B'/'C' → 只重生该套;regen_role 配合 → 只重生该套的该张; custom_prompt → 人工追加提示词。重生模式跳过文案生成,复用已有合格文案。 """ logger.info("run_generation_pipeline start: task_id=%s", task_id) # 幂等双保险:抢不到锁说明已有 worker 在跑同一 task,直接跳过(防重复烧钱)。 if not _acquire_run_lock(task_id): logger.warning("run_generation_pipeline 跳过(已有实例在跑): task_id=%s", task_id) return {"task_id": task_id, "status": "skipped_locked"} db = _get_db() seq = 0 workspace_id = 0 # G3坑修复:确保 except 里能用真实 workspace_id try: from app.constants.enums import TaskStatus from app.workers.pipeline_steps import ( load_task_and_product, decrypt_user_key, build_clients_and_clear_key, build_product_dict, load_flywheel_context, ) from app.workers.pipeline_io import run_text_generation, run_image_generation from app.core.config import get_settings # Step1: 查任务 + 产品 task, product = load_task_and_product(db, task_id) if not task: return {"task_id": task_id, "status": "not_found"} workspace_id = task.workspace_id # Step2: Fernet 解密(局部变量,不传出) plain_key = decrypt_user_key(db, task.operator_id, workspace_id) # Step3: 构建 AIClients clients = build_clients_and_clear_key(plain_key) plain_key = None # 用完即清除,基石B task.status = TaskStatus.GENERATING db.commit() # Step4: 推 task_started + 飞轮上下文 seq += 1 _push_event_sync(task_id, workspace_id, "task_started", { "task_id": task_id, "total_text": task.text_count, "total_image": task.image_count * 3, }, seq) product_dict = build_product_dict(product) flywheel_fragment, flywheel_ctx = load_flywheel_context(db, workspace_id, task.product_id, product_dict) if flywheel_fragment: seq += 1 _push_event_sync(task_id, workspace_id, "flywheel_injected", { "recent_preference": flywheel_ctx.get("recent_preference", ""), "reject_reasons": flywheel_ctx.get("reject_reasons", []), }, seq) # Step5: 文案生成(S1: 返回 needs_replenish=合格数<目标数,触发后台补充) # R2 重生模式:跳过文案重生,复用库内已有合格文案作生图依据(重生只针对图片) _is_regen = bool(regen_strategy or regen_role) if _is_regen: from app.models.task import TextCandidate as _TC _existing = db.query(_TC).filter( _TC.task_id == task_id ).order_by(_TC.id.asc()).first() if not _existing: # 重生依赖已有文案作生图语境;一条都没有时硬失败,绝不以空文案降级生图(质量崩) from app.constants.enums import TaskStatus as _TS logger.warning("重生无可复用文案,拒绝空语境生图: task_id=%s", task_id) task.status = _TS.PENDING_SELECTION db.commit() _push_event_sync(task_id, workspace_id, "error", { "code": 40002, "message": "无可复用文案,请先完成文案生成再重生图片", }, seq + 1) return {"task_id": task_id, "status": "regen_no_text"} candidates_raw = [_json_loads_safe(_existing.content)] needs_replenish = False else: candidates_raw, seq, needs_replenish = run_text_generation( db, clients, task, product_dict, flywheel_fragment, _push_event_sync, workspace_id, seq, ) if needs_replenish: # 合格文案不足用户目标条数:先展示已合格的,后台异步补充(先展示后台补铁律) # 注:candidates_raw 是本轮原始生成数(含被过滤的低分条),非合格入库数; # 真实合格数以库内 saved_count 为准,补充子任务会按库内实际缺口补到够。 logger.warning( "文案合格数不足将后台补充: task_id=%s 本轮生成=%d 目标=%d", task_id, len(candidates_raw), task.text_count, ) try: replenish_text_candidates.delay(task_id) except Exception as _re: logger.error("触发后台补充失败: task_id=%s err=%s", task_id, _re) # Step6+7+8: 图片生成 + 后处理 + 存盘 first_copy = candidates_raw[0] if candidates_raw else {} settings = get_settings() seq = run_image_generation( db, clients, task, product_dict, _push_event_sync, workspace_id, seq, first_copy, settings.UPLOAD_BASE_PATH, regen_strategy=regen_strategy, regen_role=regen_role, custom_prompt=custom_prompt, ) # 最终状态 + task_done task.status = TaskStatus.PENDING_SELECTION db.commit() seq += 1 _push_event_sync(task_id, workspace_id, "task_done", { "task_id": task_id, "status": "pending_selection" }, seq) logger.info("run_generation_pipeline done: task_id=%s", task_id) return {"task_id": task_id, "status": "pending_selection"} except Exception as exc: logger.error("run_generation_pipeline failed: task_id=%s err=%s", task_id, exc) # 重试耗尽才落 FAILED 终态;还能重试则保持 GENERATING(别伪装成"待开始")。 # B6:失败要可见,区分"没跑(pending)"和"跑挂了(failed)"。 exhausted = self.request.retries >= self.max_retries try: from app.models.task import GenerationTask as GT from app.constants.enums import TaskStatus as TS t = db.query(GT).filter(GT.id == task_id).first() if t: t.status = TS.FAILED if exhausted else TS.GENERATING db.commit() except Exception: pass if exhausted: # 彻底失败才推 error(前端弹 FatalErrorBanner);重试中不推,避免误报"生成失败"。 _push_event_sync(task_id, workspace_id, "error", {"code": 50001, "message": str(exc)}, seq + 1) return {"task_id": task_id, "status": "failed", "error": str(exc)} raise self.retry(exc=exc) finally: # 释放幂等锁:本次执行结束(成功/失败/重试间隙)即放锁。 # retry 是串行的(旧执行抛异常结束→30s后新执行重新抢锁),非并发,故每次都放。 # 真正的并发威胁(broker 窗口内重投)由 visibility_timeout=2h 兜底。 _release_run_lock(task_id) db.close() @celery_app.task( bind=True, name="app.workers.tasks.run_fission_pipeline", max_retries=2, default_retry_delay=30, queue="generation", ) def run_fission_pipeline(self, fission_id: int, source_task_id: int) -> dict: """第11环裂变主任务:一次 LLM 出 N 套完整笔记包(对齐上线版 split.js)。 只接收 id(基石B),不接收 key。幂等锁防重复烧钱。 解密 key→构建 clients→委托 fission_pipeline.execute_fission_pipeline。 """ logger.info("run_fission_pipeline start: fission_id=%s src=%s", fission_id, source_task_id) lock_key = f"fission:lock:{fission_id}" import redis as _redis from app.core.config import get_settings as _gs _r = _redis.from_url(_gs().REDIS_URL, decode_responses=True) if not _r.set(lock_key, "1", nx=True, ex=7500): logger.warning("run_fission_pipeline 跳过(已在跑): fission_id=%s", fission_id) return {"fission_id": fission_id, "status": "skipped_locked"} db = _get_db() try: from app.models.task import GenerationTask from app.workers.pipeline_steps import decrypt_user_key, build_clients_and_clear_key from app.services.fission_pipeline import execute_fission_pipeline src = db.query(GenerationTask).filter(GenerationTask.id == source_task_id).first() if not src: return {"fission_id": fission_id, "status": "not_found"} plain_key = decrypt_user_key(db, src.operator_id, src.workspace_id) clients = build_clients_and_clear_key(plain_key) plain_key = None # 用完即清,基石B return execute_fission_pipeline(db, clients, fission_id, source_task_id) except Exception as exc: logger.error("run_fission_pipeline failed: fission_id=%s err=%s", fission_id, exc) exhausted = self.request.retries >= self.max_retries if exhausted: try: from app.models.fission import FissionTask ft = db.query(FissionTask).filter(FissionTask.id == fission_id).first() if ft: ft.status = "failed"; db.commit() except Exception: pass return {"fission_id": fission_id, "status": "failed", "error": str(exc)} raise self.retry(exc=exc) finally: _r.delete(lock_key) db.close()