541 lines
21 KiB
Python
541 lines
21 KiB
Python
"""DSH content-production worker.
|
|
|
|
The workbench API owns configuration, validation, queuing, human review and
|
|
delivery. This process leases ``campaign-*`` jobs from the shared SQLite queue
|
|
and executes the existing content generation pipeline.
|
|
|
|
Run:
|
|
YAWAN_EXECUTION_BACKEND=dsh python -m app.dsh_worker
|
|
|
|
The worker must share the same YAWAN_DATA_DIR as the API process (SQLite,
|
|
knowledge bases, vector indexes and artifacts directory).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import signal
|
|
import socket
|
|
import threading
|
|
import time
|
|
import uuid
|
|
|
|
from sqlmodel import Session, select
|
|
|
|
from .agents import runner
|
|
from . import campaign_jobs
|
|
from . import planner
|
|
from .db import engine, init_db
|
|
from .models import (
|
|
Campaign,
|
|
CampaignRun,
|
|
AgentSession,
|
|
Notification,
|
|
RunEvent,
|
|
Segment,
|
|
TaskQueueItem,
|
|
now,
|
|
)
|
|
from . import progress
|
|
from . import task_queue
|
|
|
|
logger = logging.getLogger("yawen.dsh")
|
|
|
|
STOP = threading.Event()
|
|
WORKER_ID = os.getenv("YAWAN_DSH_WORKER_ID") or f"dsh-{socket.gethostname()}-{os.getpid()}-{uuid.uuid4().hex[:6]}"
|
|
POLL_SECONDS = float(os.getenv("YAWAN_DSH_POLL_SECONDS", "2.0"))
|
|
HEARTBEAT_SECONDS = int(os.getenv("YAWAN_DSH_HEARTBEAT_SECONDS", "60"))
|
|
|
|
|
|
def _record_event(session: Session, run: CampaignRun, event_type: str,
|
|
status: str = "", payload: dict | None = None) -> None:
|
|
session.add(RunEvent(
|
|
run_id=run.id, campaign_id=run.campaign_id, trace_id=run.trace_id,
|
|
event_type=event_type, status=status, payload=payload or {},
|
|
))
|
|
|
|
|
|
def _set_campaign_review(session: Session, campaign_id: int, has_failure: bool | None = None) -> None:
|
|
campaign = session.get(Campaign, campaign_id)
|
|
if not campaign or campaign.status == "done":
|
|
return
|
|
segments = session.exec(select(Segment).where(Segment.campaign_id == campaign_id)).all()
|
|
campaign.has_failure = has_failure if has_failure is not None else any(x.status == "failed" for x in segments)
|
|
if all(x.status in ("approved", "skipped") for x in segments):
|
|
campaign.status = "composing_ready"
|
|
else:
|
|
campaign.status = "human_review"
|
|
session.add(campaign)
|
|
|
|
|
|
def _recover_orphan_run(session: Session, run: CampaignRun, payload: dict) -> None:
|
|
"""Reset stale runtime state before the leased worker owns the run.
|
|
|
|
A lease is only stolen after heartbeat timeout, so any running session/segment
|
|
here belongs to a dead worker. The same applies to an immediate in-process
|
|
retry: the previous job thread is already gone.
|
|
"""
|
|
for runtime_session in session.exec(select(AgentSession).where(
|
|
AgentSession.run_id == run.id,
|
|
AgentSession.status == "running")).all():
|
|
runtime_session.status = "interrupted"
|
|
runtime_session.updated_at = now()
|
|
session.add(runtime_session)
|
|
|
|
segment_id = int(payload.get("segment_id") or 0)
|
|
if segment_id:
|
|
seg = session.get(Segment, segment_id)
|
|
if seg and seg.campaign_id == run.campaign_id and seg.status == "running":
|
|
seg.status = "pending"
|
|
session.add(seg)
|
|
else:
|
|
# 全量任务的上一次尝试可能已把 run 置为 failed；只要队列重新租约，
|
|
# 所有遗留 running 段都属于上一个已结束 worker，应恢复为可执行状态。
|
|
for seg in session.exec(select(Segment).where(
|
|
Segment.campaign_id == run.campaign_id,
|
|
Segment.status == "running")).all():
|
|
seg.status = "pending"
|
|
session.add(seg)
|
|
session.commit()
|
|
|
|
|
|
def _run_full(session: Session, run: CampaignRun) -> str:
|
|
campaign = session.get(Campaign, run.campaign_id)
|
|
if run.cancel_requested:
|
|
run.status = "cancelled"
|
|
run.finished_at = now()
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
session.add(campaign)
|
|
session.add(run)
|
|
_record_event(session, run, "run_finished", "cancelled", {"reason": "cancel_before_start"})
|
|
session.commit()
|
|
return run.status
|
|
if run.pause_requested:
|
|
run.status = "paused"
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
session.add(campaign)
|
|
session.add(run)
|
|
_record_event(session, run, "run_finished", "paused", {"reason": "pause_before_start"})
|
|
session.commit()
|
|
return run.status
|
|
run.status = "running"
|
|
run.started_at = run.started_at or now()
|
|
run.trace_id = run.trace_id or uuid.uuid4().hex
|
|
run.checkpoint = {**(run.checkpoint or {}), "phase": "running", "worker_id": WORKER_ID}
|
|
session.add(run)
|
|
_record_event(session, run, "run_started", "running", {"worker_id": WORKER_ID})
|
|
session.commit()
|
|
|
|
result = runner.run_campaign(session, run.campaign_id, run_id=run.id)
|
|
session.refresh(run)
|
|
campaign = session.get(Campaign, run.campaign_id)
|
|
|
|
if run.cancel_requested:
|
|
run.status = "cancelled"
|
|
run.finished_at = now()
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
elif str(result.get("status")) == "paused" or run.pause_requested:
|
|
run.status = "paused"
|
|
elif str(result.get("status")) in ("failed", "error") or (campaign and campaign.has_failure):
|
|
run.status = "failed"
|
|
run.error = str(result.get("error") or "")[:500]
|
|
run.finished_at = now()
|
|
else:
|
|
run.status = "awaiting_review"
|
|
run.finished_at = now()
|
|
|
|
run.checkpoint = {
|
|
**(run.checkpoint or {}),
|
|
"phase": "finished",
|
|
"campaign_status": result.get("status"),
|
|
"worker_id": WORKER_ID,
|
|
"results": result.get("results", [])[-100:],
|
|
}
|
|
session.add(run)
|
|
if campaign:
|
|
session.add(campaign)
|
|
_record_event(session, run, "run_finished", run.status, {"campaign_status": result.get("status")})
|
|
session.commit()
|
|
return run.status
|
|
|
|
|
|
def _run_step(session: Session, run: CampaignRun, segment_id: int) -> str:
|
|
campaign = session.get(Campaign, run.campaign_id)
|
|
if run.cancel_requested:
|
|
run.status = "cancelled"
|
|
run.finished_at = now()
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
session.add(campaign)
|
|
session.add(run)
|
|
_record_event(session, run, "run_finished", "cancelled", {"reason": "cancel_before_start"})
|
|
session.commit()
|
|
return run.status
|
|
if run.pause_requested:
|
|
run.status = "paused"
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
session.add(campaign)
|
|
session.add(run)
|
|
_record_event(session, run, "run_finished", "paused", {"reason": "pause_before_start"})
|
|
session.commit()
|
|
return run.status
|
|
segment = session.get(Segment, segment_id)
|
|
if not campaign or not segment or segment.campaign_id != campaign.id:
|
|
raise ValueError(f"逐段任务无效：run={run.id}, segment={segment_id}")
|
|
if segment.status not in ("pending", "rework", "failed"):
|
|
run.status = "completed"
|
|
run.error = ""
|
|
run.finished_at = now()
|
|
run.checkpoint = {**(run.checkpoint or {}), "phase": "skipped", "reason": segment.status}
|
|
session.add(run)
|
|
session.commit()
|
|
return run.status
|
|
|
|
run.status = "running"
|
|
run.started_at = run.started_at or now()
|
|
run.trace_id = run.trace_id or uuid.uuid4().hex
|
|
run.checkpoint = {**(run.checkpoint or {}), "phase": "running", "mode": "step",
|
|
"segment_ids": [segment.id], "worker_id": WORKER_ID}
|
|
session.add(run)
|
|
_record_event(session, run, "run_started", "running", {"mode": "step", "segment_id": segment.id})
|
|
session.commit()
|
|
|
|
progress.publish(int(campaign.owner_user_id or 0), "run_step_start",
|
|
campaign_id=campaign.id, segment_id=segment.id,
|
|
title=segment.title, backend="dsh")
|
|
try:
|
|
result = runner.execute_segment(session, segment, campaign, run_id=run.id)
|
|
if runner._supervisor_gate(session, segment, campaign, run.id) == "redo":
|
|
segment = session.get(Segment, segment.id)
|
|
result = runner.execute_segment(session, segment, campaign, run_id=run.id)
|
|
status = str(result.get("status") or "failed")
|
|
run.checkpoint = {**(run.checkpoint or {}), "last_result": status}
|
|
if status == "waiting_approval":
|
|
run.status = "awaiting_review"
|
|
elif status in ("human_review", "rework", "approved"):
|
|
run.status = "completed"
|
|
run.finished_at = now()
|
|
else:
|
|
run.status = "failed"
|
|
run.error = str(result.get("error") or "")[:500]
|
|
run.finished_at = now()
|
|
_set_campaign_review(session, campaign.id)
|
|
session.add(run)
|
|
_record_event(session, run, "run_finished", run.status, {"segment_status": status})
|
|
session.commit()
|
|
return run.status
|
|
except Exception as exc:
|
|
session.rollback()
|
|
segment = session.get(Segment, segment_id)
|
|
run = session.get(CampaignRun, run.id)
|
|
campaign = session.get(Campaign, run.campaign_id)
|
|
if segment:
|
|
segment.status = "failed"
|
|
segment.failure = {"type": "dsh_step_failed", "detail": str(exc)[:300]}
|
|
session.add(segment)
|
|
run.status = "failed"
|
|
run.error = str(exc)[:500]
|
|
run.finished_at = now()
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
session.add(campaign)
|
|
session.add(run)
|
|
session.commit()
|
|
raise
|
|
|
|
|
|
def _segment_redo(session: Session, payload: dict) -> None:
|
|
segment_id = int(payload.get("segment_id") or 0)
|
|
opinion = str(payload.get("opinion") or "")
|
|
expected_agent = int(payload.get("expected_agent_id") or 0)
|
|
expected_round = int(payload.get("expected_round") or 0)
|
|
seg = session.get(Segment, segment_id)
|
|
campaign = session.get(Campaign, seg.campaign_id) if seg else None
|
|
if not seg or not campaign:
|
|
raise ValueError("重做任务对应的内容段不存在")
|
|
if seg.status == "running":
|
|
# 队列租约过期后才可能重新领取到该确定性任务；旧 Worker 已被视为死亡。
|
|
seg.status = "rework"
|
|
session.add(seg)
|
|
session.commit()
|
|
if seg.status not in ("rework", "failed"):
|
|
return
|
|
if expected_agent and int(seg.agent_id or 0) != expected_agent:
|
|
return
|
|
if expected_round and int(seg.round or 0) != expected_round:
|
|
return
|
|
seg.status = "running"
|
|
session.add(seg)
|
|
session.commit()
|
|
progress.publish(int(campaign.owner_user_id or 0), "redo_start", campaign_id=campaign.id,
|
|
segment_id=seg.id, round=seg.round, title=seg.title,
|
|
reason="user", backend="dsh")
|
|
try:
|
|
runner.execute_segment(session, seg, campaign, auto_retry_qc=False, qc_issues=opinion)
|
|
seg2 = session.get(Segment, segment_id)
|
|
progress.publish(int(campaign.owner_user_id or 0), "redo_done", campaign_id=campaign.id,
|
|
segment_id=segment_id, status=(seg2.status if seg2 else "?"),
|
|
round=(seg2.round if seg2 else 0), title=(seg2.title if seg2 else ""))
|
|
_set_campaign_review(session, campaign.id)
|
|
session.commit()
|
|
except Exception as exc:
|
|
session.rollback()
|
|
seg = session.get(Segment, segment_id)
|
|
campaign = session.get(Campaign, campaign.id)
|
|
if seg:
|
|
seg.status = "failed"
|
|
seg.failure = {"type": "dsh_redo_failed", "detail": str(exc)[:300]}
|
|
session.add(seg)
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
campaign.has_failure = True
|
|
session.add(campaign)
|
|
session.commit()
|
|
raise
|
|
|
|
|
|
def _segment_fix(session: Session, payload: dict) -> None:
|
|
segment_id = int(payload.get("segment_id") or 0)
|
|
seg = session.get(Segment, segment_id)
|
|
if seg and seg.status in ("skipped", "approved"):
|
|
# skip/disable 的管理动作已在 API 请求中完成；approved 说明人工已接管，旧任务作废。
|
|
return
|
|
if seg and seg.status == "running":
|
|
# 上一执行进程租约超时，遗留 running 段恢复为 pending 后重试。
|
|
seg.status = "pending"
|
|
seg.failure = {}
|
|
session.add(seg)
|
|
session.commit()
|
|
if seg and seg.status not in ("pending", "failed"):
|
|
return
|
|
fix = payload.get("fix") if isinstance(payload.get("fix"), dict) else {}
|
|
try:
|
|
result = runner.fix_segment(session, segment_id, fix, execute=True, apply_changes=False)
|
|
if result.get("status") == "error":
|
|
raise ValueError(str(result.get("error") or "修复执行失败"))
|
|
seg = session.get(Segment, segment_id)
|
|
if seg:
|
|
_set_campaign_review(session, seg.campaign_id)
|
|
session.commit()
|
|
except Exception as exc:
|
|
session.rollback()
|
|
seg = session.get(Segment, segment_id)
|
|
if seg:
|
|
campaign = session.get(Campaign, seg.campaign_id)
|
|
seg.status = "failed"
|
|
seg.failure = {"type": "dsh_fix_failed", "detail": str(exc)[:300]}
|
|
session.add(seg)
|
|
if campaign:
|
|
campaign.status = "human_review"
|
|
campaign.has_failure = True
|
|
session.add(campaign)
|
|
session.commit()
|
|
raise
|
|
|
|
|
|
def _campaign_parse(session: Session, payload: dict) -> None:
|
|
campaign_id = int(payload.get("campaign_id") or 0)
|
|
campaign_jobs.run_parse_job(session, campaign_id, str(payload.get("advice") or ""))
|
|
|
|
|
|
def _campaign_plan(session: Session, payload: dict) -> None:
|
|
campaign_id = int(payload.get("campaign_id") or 0)
|
|
template_id = int(payload.get("template_id") or 0)
|
|
result = planner.run_plan_job(
|
|
session,
|
|
campaign_id,
|
|
template_id=template_id,
|
|
advice=str(payload.get("advice") or ""),
|
|
created_by=int(payload.get("created_by") or 0),
|
|
)
|
|
if result.get("status") == "failed":
|
|
logger.warning("DSH plan business failure: %s", str(result.get("error"))[:200])
|
|
|
|
|
|
def _campaign_compose(session: Session, payload: dict) -> None:
|
|
campaign_id = int(payload.get("campaign_id") or 0)
|
|
result = campaign_jobs.run_compose_job(session, campaign_id)
|
|
if result.get("status") == "error":
|
|
# 与 inline 同步接口一致：业务校验失败不做队列重试，工作台保留错误并允许用户重试点。
|
|
logger.warning("DSH compose business failure: %s", str(result.get("error"))[:200])
|
|
|
|
|
|
def process_job(job: TaskQueueItem) -> None:
|
|
with Session(engine) as session:
|
|
payload = dict(job.payload or {})
|
|
key = job.task_key or ""
|
|
logger.info("DSH lease job %s: %s", job.id, key)
|
|
|
|
if key.startswith("campaign-run:"):
|
|
run_id = int(payload.get("run_id") or key.rsplit(":", 1)[-1])
|
|
run = session.get(CampaignRun, run_id)
|
|
if not run:
|
|
raise ValueError(f"运行不存在：{key}")
|
|
_recover_orphan_run(session, run, payload)
|
|
session.refresh(run)
|
|
try:
|
|
if run.status in ("completed", "cancelled"):
|
|
return
|
|
if payload.get("mode") == "step" or payload.get("segment_id"):
|
|
_run_step(session, run, int(payload.get("segment_id") or 0))
|
|
else:
|
|
_run_full(session, run)
|
|
finally:
|
|
# 预算池是进程内对象；DSH 每次执行结束都必须释放，防止长驻 Worker 内存泄漏。
|
|
from .agents.runtime import release_run_budget
|
|
release_run_budget(run_id)
|
|
return
|
|
|
|
if key.startswith("campaign-segment-redo:"):
|
|
_segment_redo(session, payload)
|
|
return
|
|
|
|
if key.startswith("campaign-segment-fix:"):
|
|
_segment_fix(session, payload)
|
|
return
|
|
|
|
if key.startswith("campaign-plan:"):
|
|
_campaign_plan(session, payload)
|
|
return
|
|
|
|
if key.startswith("campaign-parse:"):
|
|
_campaign_parse(session, payload)
|
|
return
|
|
|
|
if key.startswith("campaign-compose:"):
|
|
_campaign_compose(session, payload)
|
|
return
|
|
|
|
raise ValueError(f"DSH 不识别的内容生产任务：{key}")
|
|
|
|
|
|
def _mark_dead_job_state(session: Session, job: TaskQueueItem, error: str) -> None:
|
|
"""Keep business state visible when all queue retry attempts are exhausted."""
|
|
payload = dict(job.payload or {})
|
|
key = job.task_key or ""
|
|
campaign_id = int(payload.get("campaign_id") or 0)
|
|
run_id = int(payload.get("run_id") or 0)
|
|
detail = str(error)[:300]
|
|
|
|
if key.startswith("campaign-run:") or run_id:
|
|
run = session.get(CampaignRun, run_id) if run_id else None
|
|
if run:
|
|
campaign_id = campaign_id or run.campaign_id
|
|
run.status = "failed"
|
|
run.error = detail
|
|
run.finished_at = now()
|
|
session.add(run)
|
|
elif key.startswith("campaign-plan:") or key.startswith("campaign-parse:"):
|
|
campaign = session.get(Campaign, campaign_id)
|
|
if campaign:
|
|
campaign.status = "plan_failed" if key.startswith("campaign-plan:") else "parse_failed"
|
|
notes = list(campaign.notes or [])
|
|
label = "规划" if key.startswith("campaign-plan:") else "拆解"
|
|
notes.append(f"DSH {label}任务进入死信队列：{detail[:180]}")
|
|
campaign.notes = notes[-20:]
|
|
session.add(campaign)
|
|
elif key.startswith("campaign-compose:"):
|
|
campaign = session.get(Campaign, campaign_id)
|
|
if campaign and campaign.status != "done":
|
|
campaign.status = "composing_ready"
|
|
session.add(campaign)
|
|
|
|
segment_id = int(payload.get("segment_id") or 0)
|
|
if segment_id:
|
|
seg = session.get(Segment, segment_id)
|
|
if seg and seg.status == "running":
|
|
seg.status = "failed"
|
|
seg.failure = {"type": "dsh_job_dead", "detail": detail}
|
|
session.add(seg)
|
|
campaign_id = campaign_id or seg.campaign_id
|
|
|
|
campaign = session.get(Campaign, campaign_id) if campaign_id else None
|
|
if campaign and campaign.status not in ("done", "parse_failed", "plan_failed", "composing_ready"):
|
|
campaign.status = "human_review"
|
|
campaign.has_failure = True
|
|
session.add(campaign)
|
|
session.commit()
|
|
|
|
|
|
def _heartbeat(job_id: int) -> threading.Thread:
|
|
def loop():
|
|
while not STOP.wait(HEARTBEAT_SECONDS):
|
|
try:
|
|
with Session(engine) as session:
|
|
ok = task_queue.heartbeat(session, TaskQueueItem, job_id, WORKER_ID)
|
|
if not ok:
|
|
return
|
|
except Exception:
|
|
logger.exception("DSH heartbeat failed: job=%s", job_id)
|
|
thread = threading.Thread(target=loop, name=f"dsh-heartbeat-{job_id}", daemon=True)
|
|
thread.start()
|
|
return thread
|
|
|
|
|
|
def worker_loop() -> None:
|
|
init_db()
|
|
logger.info("DSH worker started: %s (poll=%ss, heartbeat=%ss)", WORKER_ID, POLL_SECONDS, HEARTBEAT_SECONDS)
|
|
while not STOP.is_set():
|
|
job = None
|
|
try:
|
|
with Session(engine) as session:
|
|
job = task_queue.lease_next(session, TaskQueueItem, WORKER_ID,
|
|
task_key_prefix="campaign-%")
|
|
if job is None:
|
|
STOP.wait(POLL_SECONDS)
|
|
continue
|
|
|
|
beat = _heartbeat(job.id)
|
|
process_job(job)
|
|
beat.join(timeout=1)
|
|
with Session(engine) as session:
|
|
task_queue.complete(session, TaskQueueItem, job.id, WORKER_ID)
|
|
logger.info("DSH job completed: %s", job.id)
|
|
except Exception as exc:
|
|
logger.exception("DSH job failed: %s", getattr(job, "id", "?"))
|
|
try:
|
|
with Session(engine) as session:
|
|
if job is not None:
|
|
status = task_queue.fail(session, TaskQueueItem, job.id, WORKER_ID, str(exc))
|
|
if status == "dead":
|
|
_mark_dead_job_state(session, job, str(exc))
|
|
payload = dict(job.payload or {})
|
|
campaign_id = int(payload.get("campaign_id") or 0)
|
|
if not campaign_id and payload.get("run_id"):
|
|
run = session.get(CampaignRun, int(payload["run_id"]))
|
|
campaign_id = run.campaign_id if run else 0
|
|
if campaign_id:
|
|
session.add(Notification(
|
|
type="failure", ref_id=campaign_id,
|
|
message=f"DSH 内容生产任务进入死信队列：{str(exc)[:160]}",
|
|
))
|
|
session.commit()
|
|
except Exception:
|
|
logger.exception("Failed to mark DSH job failed: %s", getattr(job, "id", "?"))
|
|
STOP.wait(min(30.0, POLL_SECONDS * 3))
|
|
logger.info("DSH worker stopped: %s", WORKER_ID)
|
|
|
|
|
|
def _signal_handler(signum, frame): # pragma: no cover
|
|
logger.info("DSH received signal %s, stopping...", signum)
|
|
STOP.set()
|
|
|
|
|
|
def main() -> None:
|
|
logging.basicConfig(
|
|
level=os.getenv("YAWAN_DSH_LOG_LEVEL", "INFO").upper(),
|
|
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
|
)
|
|
if hasattr(signal, "SIGTERM"):
|
|
signal.signal(signal.SIGTERM, _signal_handler)
|
|
signal.signal(signal.SIGINT, _signal_handler)
|
|
worker_loop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|