Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 25 additions & 10 deletions apps/worker/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -234,27 +234,41 @@ def run_worker() -> None:
│ 闲时自动处理 │ 全天检测 │ 全天检测 │
└─────────────────────────────────────────────────────────┘
"""
# High 3e:显式配置 max_instances / misfire_grace_time / coalesce,避免
# 重复触发与 misfire 丢失;用 apscheduler 的 ThreadPoolExecutor(max_workers=3)
# 替代默认单线程池,允许 topic_dispatch / cs_feed / brief 适度并发
from apscheduler.executors.pool import ThreadPoolExecutor as APSThreadPoolExecutor

scheduler = BlockingScheduler(timezone="UTC")
scheduler.add_executor(APSThreadPoolExecutor(max_workers=3))

# 公共 job 配置:单实例 + 5 分钟 misfire 容忍 + 合并错过的触发
_job_kwargs = {
"max_instances": 1,
"misfire_grace_time": 300,
"coalesce": True,
"replace_existing": True,
}

settings = get_settings()

# 每整点检查主题调度(UTC 时间)
# 每整点检查主题调度(UTC 时间)—— 整点第 0 分钟
scheduler.add_job(
topic_dispatch_job,
trigger=CronTrigger(minute=0),
id="topic_dispatch",
replace_existing=True,
**_job_kwargs,
)
logger.info("✅ 已添加:主题分发任务(每小时整点,UTC)")

# CS 分类订阅调度(每小时整点)
# CS 分类订阅调度 —— 错开 5 分钟,避免与 topic_dispatch 同分钟抢线程
scheduler.add_job(
cs_feed_dispatch_job,
trigger=CronTrigger(minute=0),
trigger=CronTrigger(minute=5),
id="cs_feed_dispatch",
replace_existing=True,
**_job_kwargs,
)
logger.info("✅ 已添加:CS分类订阅调度任务(每小时整点,UTC)")
logger.info("✅ 已添加:CS分类订阅调度任务(每小时 :05,UTC)")

# 每日简报(从数据库读取 cron 表达式)
from packages.storage.db import session_scope
Expand All @@ -273,7 +287,7 @@ def run_worker() -> None:
brief_job,
trigger=daily_trigger,
id="daily_brief",
replace_existing=True,
**_job_kwargs,
)
logger.info(
"✅ 已添加:每日简报任务(cron: %s)",
Expand All @@ -286,16 +300,17 @@ def run_worker() -> None:
weekly_graph_job,
trigger=weekly_trigger,
id="weekly_graph",
replace_existing=True,
**_job_kwargs,
)
logger.info("✅ 已添加:每周图谱维护任务(UTC 周日 22:00)")

# 优雅关闭
# 优雅关闭(High 3f:等待进行中任务跑完,避免已下载 PDF 未 set_pdf_path
# 的中间态丢失;wait=True + 60s 超时兜底)
def _graceful_stop(*_: object) -> None:
logger.info("收到终止信号,正在关闭...")
stop_event.set()
stop_idle_processor() # 停止闲时处理器
scheduler.shutdown(wait=False)
scheduler.shutdown(wait=True)
logger.info("Worker 已关闭")

signal.signal(signal.SIGINT, _graceful_stop)
Expand Down
207 changes: 139 additions & 68 deletions packages/ai/cs_feed_orchestrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,83 +68,154 @@ def sync_categories(self):
session.close()

def run(self):
"""每小时执行一次"""
session = SessionLocal()
"""每小时执行一次(High 3a:每 sub 独立 session,异常隔离)"""
# 先读订阅列表(独立 session,读后即关)
read_session = SessionLocal()
try:
repo = CSFeedRepository(session)
repo = CSFeedRepository(read_session)
subs = repo.get_active_subscriptions()
now = datetime.now(UTC)
digest: list[tuple[str, int, list[str]]] = []

for sub in subs:
# 冷却中检查
if sub.status == "cool_down" and sub.cool_down_until and now < sub.cool_down_until:
logger.info(
"[CSFeed] Skipping %s (cool down until %s)",
sub.category_code,
sub.cool_down_until,
)
continue
sub_specs = [
(
s.category_code,
s.status,
s.cool_down_until,
s.last_run_at,
s.last_run_count,
s.daily_limit,
)
for s in subs
]
finally:
read_session.close()

# 每日配额检查
today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
if sub.last_run_at and sub.last_run_at >= today_start:
remaining = sub.daily_limit - sub.last_run_count
else:
remaining = sub.daily_limit
now = datetime.now(UTC)
digest: list[tuple[str, int, list[str]]] = []

if remaining <= 0:
logger.info("[CSFeed] Skipping %s (daily limit reached)", sub.category_code)
continue
for (
category_code,
status,
cool_down_until,
last_run_at,
last_run_count,
daily_limit,
) in sub_specs:
# 冷却中检查
if status == "cool_down" and cool_down_until and now < cool_down_until:
logger.info(
"[CSFeed] Skipping %s (cool down until %s)",
category_code,
cool_down_until,
)
continue

# 请求间隔
if not self.bucket.acquire(timeout=30):
logger.warning("[CSFeed] Token bucket timeout, skipping %s", sub.category_code)
continue
time.sleep(REQUEST_INTERVAL)
# 每日配额检查
today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
if last_run_at and last_run_at >= today_start:
remaining = daily_limit - last_run_count
else:
remaining = daily_limit

# 抓取
try:
client = ArxivClient()
papers = client.fetch_latest(
query=f"cat:{sub.category_code}",
max_results=remaining,
days_back=7,
)
from packages.storage.repositories import PaperRepository

paper_repo = PaperRepository(session)
count = 0
titles: list[str] = []
for p in papers:
paper_repo.upsert_paper(p)
count += 1
title = getattr(p, "title", None)
if title:
titles.append(title)

repo.update_run_status(sub.category_code, count)
logger.info("[CSFeed] %s: ingested %d papers", sub.category_code, count)
if count > 0:
digest.append((sub.category_code, count, titles))
if remaining <= 0:
logger.info("[CSFeed] Skipping %s (daily limit reached)", category_code)
continue

except Exception as e:
err_str = str(e)
if "429" in err_str or "Too Many Requests" in err_str:
repo.set_cool_down(
sub.category_code, now + timedelta(minutes=COOL_DOWN_MINUTES)
)
logger.warning(
"[CSFeed] Rate limited %s, cool down 30min", sub.category_code
# 请求间隔
if not self.bucket.acquire(timeout=30):
logger.warning("[CSFeed] Token bucket timeout, skipping %s", category_code)
continue
time.sleep(REQUEST_INTERVAL)

# 每 sub 独立 session:一个 sub 抓取异常不会污染下个 sub 的脏数据
sub_session = SessionLocal()
try:
sub_repo = CSFeedRepository(sub_session)
client = ArxivClient()
papers = client.fetch_latest(
query=f"cat:{category_code}",
max_results=remaining,
days_back=7,
)
from packages.storage.repositories import PaperRepository

paper_repo = PaperRepository(sub_session)
count = 0
titles: list[str] = []
paper_ids: list[str] = []
for p in papers:
saved = paper_repo.upsert_paper(p)
count += 1
paper_ids.append(saved.id)
title = getattr(p, "title", None)
if title:
titles.append(title)

sub_repo.update_run_status(category_code, count)
sub_session.commit()
logger.info("[CSFeed] %s: ingested %d papers", category_code, count)
if count > 0:
digest.append((category_code, count, titles))
# High 3b:抓取的论文触发 embed + skim(复用 PaperPipelines),
# 此前 cs_feed 论文处于 unread 无 embedding 无 topic,只能靠
# idle_processor 事后补——改为抓取即处理
self._process_cs_papers(paper_ids)
except Exception as e:
sub_session.rollback()
err_str = str(e)
if "429" in err_str or "Too Many Requests" in err_str:
# 冷却设置需独立 session(当前已回滚)
cool_session = SessionLocal()
try:
CSFeedRepository(cool_session).set_cool_down(
category_code, now + timedelta(minutes=COOL_DOWN_MINUTES)
)
else:
logger.error("[CSFeed] Error fetching %s: %s", sub.category_code, e)
cool_session.commit()
finally:
cool_session.close()
logger.warning("[CSFeed] Rate limited %s, cool down 30min", category_code)
else:
logger.error("[CSFeed] Error fetching %s: %s", category_code, e)
finally:
sub_session.close()

# 入库后推送邮件摘要(README 宣传的「邮件推送」);SMTP 未配置时静默跳过
if digest:
self._notify_digest(digest)
finally:
session.close()
# 入库后推送邮件摘要(README 宣传的「邮件推送」);SMTP 未配置时静默跳过
if digest:
self._notify_digest(digest)

def _process_cs_papers(self, paper_ids: list[str]) -> None:
"""对 cs_feed 抓取的论文触发 embed + skim(High 3b,抓取即处理)。

复用 PaperPipelines 的 embed_paper / skim,使 cs_feed 论文不再只入库后
处于 unread 无 embedding 状态。失败不抛(仅记录日志),不阻断抓取主流程。
"""
if not paper_ids:
return
from packages.ai.pipelines import PaperPipelines
from packages.ai.rate_limiter import acquire_api, get_rate_limiter

pipelines = PaperPipelines()
limiter = get_rate_limiter()
for pid in paper_ids:
if not limiter.start_task():
logger.debug("[CSFeed] 并发满,跳过处理 %s", pid)
break
try:
if not acquire_api("embedding", timeout=30.0):
logger.warning("[CSFeed] Embedding 限流,跳过 %s", pid)
continue
try:
pipelines.embed_paper(pid)
except Exception as e:
logger.warning("[CSFeed] embed %s 失败: %s", pid, e)
continue
if not acquire_api("llm", timeout=30.0):
logger.warning("[CSFeed] LLM 限流,跳过 skim %s", pid)
continue
try:
pipelines.skim(pid)
except Exception as e:
logger.warning("[CSFeed] skim %s 失败: %s", pid, e)
finally:
limiter.end_task()

def _notify_digest(self, digest: list[tuple[str, int, list[str]]]) -> None:
"""抓取入库后发送邮件摘要;SMTP 未配置或无收件人时静默跳过"""
Expand Down
23 changes: 11 additions & 12 deletions packages/ai/pipelines/paper_pipelines.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from __future__ import annotations

import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING

if TYPE_CHECKING:
Expand Down Expand Up @@ -92,6 +92,12 @@ def _bg_auto_link(paper_ids: list[str]) -> None:
logger.warning("bg auto_link failed: %s", exc)


# High 3c:复用有界线程池替代无界 daemon 线程。此前每次 collect 都
# threading.Thread(...).start(),5 个 topic 各起线程可能失控。
# max_workers=2 限制并发 auto_link,避免线程数膨胀。
_auto_link_pool = ThreadPoolExecutor(max_workers=2, thread_name_prefix="bg-auto-link")


class PaperPipelines:
def __init__(self) -> None:
self.settings = get_settings()
Expand Down Expand Up @@ -228,11 +234,8 @@ def ingest_arxiv(

run_repo.finish(run.id)
if inserted_ids:
threading.Thread(
target=_bg_auto_link,
args=(inserted_ids,),
daemon=True,
).start()
# High 3c:提交到有界线程池,替代无界 daemon 线程
_auto_link_pool.submit(_bg_auto_link, inserted_ids)

logger.info(
"抓取完成:共 %d 篇新论文(从 %d 篇中筛选)",
Expand Down Expand Up @@ -375,12 +378,8 @@ def ingest_ieee(
topic_id=topic_id,
)

# 后台关联引用
threading.Thread(
target=_bg_auto_link,
args=(inserted_ids,),
daemon=True,
).start()
# 后台关联引用(High 3c:提交到有界线程池)
_auto_link_pool.submit(_bg_auto_link, inserted_ids)

run_repo.finish(run.id)

Expand Down
Loading