83 lines
3.6 KiB
Python
83 lines
3.6 KiB
Python
"""One bounded pool for indexing, assignment and event research; no third stage."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from sqlalchemy import select, text, update
|
|
from alphapulse.core.db import SessionLocal, engine
|
|
from alphapulse.news_research.models import (
|
|
ResearchSettings,
|
|
Package,
|
|
Event,
|
|
Attempt,
|
|
Version,
|
|
Batch, BatchEvent,
|
|
)
|
|
from alphapulse.news_research.service import process_package, process_event
|
|
from alphapulse.news_research import batches
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def run():
|
|
async with engine.connect() as lock:
|
|
if not await lock.scalar(text("SELECT pg_try_advisory_lock(20261030, 1)")):
|
|
raise RuntimeError("新闻AI worker 已运行")
|
|
await lock.commit()
|
|
async with SessionLocal() as session:
|
|
for stage in ("index", "assign"):
|
|
await session.execute(
|
|
update(Package)
|
|
.where(Package.state == stage + "_processing")
|
|
.values(state=stage + "_pending")
|
|
)
|
|
await session.execute(
|
|
update(Attempt)
|
|
.where(Attempt.state == "processing")
|
|
.values(state="failed", error="进程中断，待重试")
|
|
)
|
|
await session.execute(
|
|
update(Event).where(Event.state == "processing").values(state="pending")
|
|
)
|
|
await session.execute(
|
|
update(Version)
|
|
.where(Version.state == "processing")
|
|
.values(state="failed", error="进程中断，待重试")
|
|
)
|
|
await session.execute(update(BatchEvent).where(BatchEvent.state == 'processing').values(state='pending'))
|
|
await session.commit()
|
|
tasks = {}
|
|
try:
|
|
while True:
|
|
for task in list(tasks):
|
|
if task.done():
|
|
tasks.pop(task)
|
|
try:
|
|
task.result()
|
|
except Exception:
|
|
logger.exception("新闻研究任务失败")
|
|
async with SessionLocal() as session:
|
|
settings = await session.get(ResearchSettings, 1)
|
|
if settings and settings.enabled:
|
|
batch = await session.scalar(select(Batch).where(
|
|
Batch.phase.in_(['assigning', 'researching'])).order_by(Batch.id).limit(1))
|
|
if batch is None:
|
|
batch = await batches.freeze(session, settings)
|
|
settings = await session.get(ResearchSettings, 1, populate_existing=True, with_for_update=True)
|
|
if batch is not None and settings.enabled:
|
|
active = set(tasks.values())
|
|
actions = await batches.actions(session, batch, active,
|
|
max(0, settings.config.get('concurrency', 3) - len(tasks)))
|
|
await session.commit()
|
|
for stage, identifier in actions:
|
|
call = (process_event(identifier, settings.config, force=True)
|
|
if stage == 'event' else process_package(identifier, settings.config))
|
|
tasks[asyncio.create_task(call)] = (stage, identifier)
|
|
await asyncio.sleep(1)
|
|
finally:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO)
|
|
asyncio.run(run())
|